Skip to main content

cortiq_engine/
qtensor.rs

1//! QTensor — weight tensor with pluggable storage.
2//!
3//! Two backings, one interface:
4//! - `F32`   — owned dense floats (small models, tests). Every operation
5//!   is bit-identical to the historical `&[f32]` code paths.
6//! - `Mapped` — quantized bytes zero-copy from the CMF mmap (`q8_row` /
7//!   `q8_2f`). The matvec is fused: int8 rows × f32 activations, the
8//!   q8_2f column field folds into a pre-scale of the input
9//!   (`x'[i] = col[i]·x[i]`), so the inner loop is the same i8 dot as
10//!   q8_row. This is what lets a 15B file run in a few GB of RSS.
11//!
12//! Extension point: new dtypes = new match arm here, nothing else moves.
13
14use crate::pool::{Pool, matvec_rows, matvec_rows2};
15use cortiq_core::quant::{
16    GROUP_SIZE, Q1_TILE, Q2TP_CHUNK, Q4_TILE, Q4TP_NIB, f16_to_f32, q2tp_ladder, q2tp_sections,
17    q4tp_code, q4tp_ladder, q4tp_sections,
18};
19use cortiq_core::{CmfModel, TensorDtype};
20use std::sync::Arc;
21
22pub enum QTensor {
23    F32 {
24        data: Vec<f32>,
25        rows: usize,
26        cols: usize,
27    },
28    Mapped {
29        model: Arc<CmfModel>,
30        /// Index into the model's tensor directory.
31        idx: usize,
32        dtype: TensorDtype,
33        rows: usize,
34        cols: usize,
35        /// Per-row scales, dequantized to f32 up front (tiny).
36        row_scale: Vec<f32>,
37        /// q8_2f column field (θ), dequantized up front; empty for q8_row.
38        col_field: Vec<f32>,
39        /// Vbit only: byte offset of each row's packed data within the
40        /// tensor blob (`[rows + 1]`, computed once at load — the per-
41        /// matvec prefix scan over row bit-widths was O(rows) each call).
42        vbit_offsets: Vec<usize>,
43        /// q8-family decode repack (load-time, optional): rows in groups
44        /// of 4, interleaved in 16-byte units — one 64-byte line per
45        /// iteration feeds all 4 sdot lanes, ONE sequential weight
46        /// stream per worker instead of four (this is where llama.cpp's
47        /// repacked Q8 kernels get their bandwidth). Empty = off
48        /// (CMF_REPACK=0, non-SDOT arch, or an ineligible shape). Trades
49        /// an anonymous copy of the quants for mmap pages that go cold.
50        repack: Vec<u8>,
51    },
52}
53
54/// Load-time q8 repack gate (see `Mapped::repack`). OPT-IN
55/// (`CMF_REPACK=1`): the single-stream hypothesis LOST on Apple Silicon
56/// (M4, interleaved A/B: decode 101 vs 94 tok/s — four adjacent row
57/// streams per worker feed the prefetcher MORE memory-level parallelism
58/// than one); kept as an experiment flag for x86, where the tradeoff
59/// may land differently.
60fn repack_enabled() -> bool {
61    static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
62    *ON.get_or_init(|| {
63        std::env::var("CMF_REPACK")
64            .map(|v| v == "1")
65            .unwrap_or(cfg!(target_os = "android"))
66    })
67}
68
69/// Interleave q8 rows for the decode kernel: group g holds rows
70/// 4g..4g+4 as [r0[c], r1[c], r2[c], r3[c]] per 16-byte chunk c. Only
71/// full groups are packed — tail rows keep reading the mmap layout.
72fn q8_repack(bytes: &[u8], rows: usize, cols: usize) -> Vec<u8> {
73    #[cfg(target_arch = "aarch64")]
74    let arch_ok = sdot_enabled();
75    #[cfg(not(target_arch = "aarch64"))]
76    let arch_ok = false;
77    if !arch_ok || !repack_enabled() || rows < 256 || cols % 16 != 0 {
78        return Vec::new();
79    }
80    q8_repack_layout(bytes, rows, cols)
81}
82
83/// The pure layout transform behind `q8_repack` (tested directly —
84/// the gate depends on arch and env).
85fn q8_repack_layout(bytes: &[u8], rows: usize, cols: usize) -> Vec<u8> {
86    let groups = rows / 4;
87    let mut rep = vec![0u8; groups * 4 * cols];
88    for g in 0..groups {
89        let dst = &mut rep[g * 4 * cols..(g + 1) * 4 * cols];
90        for c in 0..cols / 16 {
91            for lane in 0..4 {
92                let src = (g * 4 + lane) * cols + c * 16;
93                dst[c * 64 + lane * 16..c * 64 + lane * 16 + 16]
94                    .copy_from_slice(&bytes[src..src + 16]);
95            }
96        }
97    }
98    rep
99}
100
101/// Prefix-sum of vbit row payload offsets (absolute within the tensor
102/// bytes). `offsets[r]..offsets[r+1]` is row r's packed data.
103fn vbit_row_offsets(bytes: &[u8], rows: usize, cols: usize) -> Vec<usize> {
104    let ng = cols / GROUP_SIZE;
105    let bits = &bytes[..rows];
106    let mut offsets = Vec::with_capacity(rows + 1);
107    let mut off = rows + rows * ng * 2;
108    for r in 0..rows {
109        offsets.push(off);
110        off += (cols * bits[r] as usize).div_ceil(8);
111    }
112    offsets.push(off);
113    offsets
114}
115
116/// `CMF_X86_BLOCKED` / `CMF_GPU_LMHEAD` / `CMF_GPU_SPLIT`, read once. They
117/// used to be read from the environment on every large matvec and on every
118/// matmat in six places — microseconds each, but also a knob that could
119/// change under a running process, which is not a thing a kernel choice
120/// should be able to do mid-sequence.
121fn blocked_enabled() -> bool {
122    use std::sync::atomic::Ordering::Relaxed;
123    match BLOCKED_OVERRIDE.load(Relaxed) {
124        1 => false,
125        2 => true,
126        _ => {
127            static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
128            *ON.get_or_init(|| {
129                std::env::var("CMF_X86_BLOCKED")
130                    .map(|v| v != "0")
131                    .unwrap_or(true)
132            })
133        }
134    }
135}
136
137static BLOCKED_OVERRIDE: std::sync::atomic::AtomicU8 = std::sync::atomic::AtomicU8::new(0);
138
139/// Force the blocked GEMM on or off, ignoring the environment; `None`
140/// restores it. For tests that need to run BOTH paths and compare them:
141/// `blocked_enabled` caches its answer for the life of the process, which
142/// is right when the environment is the only input, but leaves a test that
143/// flips `CMF_X86_BLOCKED` between two calls comparing a path against
144/// itself — or against whatever a test running in parallel latched first.
145pub fn set_blocked_override(on: Option<bool>) {
146    let v = match on {
147        None => 0,
148        Some(false) => 1,
149        Some(true) => 2,
150    };
151    BLOCKED_OVERRIDE.store(v, std::sync::atomic::Ordering::Relaxed);
152}
153
154fn gpu_lmhead_enabled() -> bool {
155    static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
156    *ON.get_or_init(|| {
157        std::env::var("CMF_GPU_LMHEAD")
158            .map(|v| v != "0")
159            .unwrap_or(true)
160    })
161}
162
163fn gpu_split_frac() -> f32 {
164    // MiMo's banked verification projects every row on the device. Plain
165    // decode must not quantize half of its O/head activation on the CPU.
166    if FULL_GPU_Q8.get() {
167        return 1.0;
168    }
169    static F: std::sync::OnceLock<f32> = std::sync::OnceLock::new();
170    *F.get_or_init(|| {
171        std::env::var("CMF_GPU_SPLIT")
172            .ok()
173            .and_then(|v| v.parse::<f32>().ok())
174            .unwrap_or(0.5)
175            .clamp(0.0, 1.0)
176    })
177}
178
179impl QTensor {
180    pub fn from_f32(data: Vec<f32>, rows: usize, cols: usize) -> Self {
181        debug_assert_eq!(data.len(), rows * cols);
182        Self::F32 { data, rows, cols }
183    }
184
185    /// Wrap a directory tensor without dequantizing the payload.
186    /// Falls back to dequantized f32 for dtypes without a fused kernel.
187    pub fn from_model(model: &Arc<CmfModel>, name: &str) -> Result<Self, String> {
188        // Indexed lookup: the linear directory scan made pipeline build
189        // O(N²) on MoE/skills files with thousands of tensors.
190        let idx = model
191            .tensor_index(name)
192            .ok_or_else(|| format!("tensor '{name}' not found in CMF directory"))?;
193        let entry = &model.tensors[idx];
194        if entry.shape.len() != 2 {
195            return Err(format!("QTensor::from_model needs 2-D, got '{name}'"));
196        }
197        let (rows, cols) = (entry.shape[0], entry.shape[1]);
198        let bytes = model.entry_bytes(entry);
199
200        match entry.dtype {
201            TensorDtype::Q8Row | TensorDtype::Q8_2f => {
202                let n = rows * cols;
203                let scales_off = n;
204                let row_scale: Vec<f32> = (0..rows)
205                    .map(|o| {
206                        f16_to_f32(u16::from_le_bytes([
207                            bytes[scales_off + o * 2],
208                            bytes[scales_off + o * 2 + 1],
209                        ]))
210                    })
211                    .collect();
212                let col_field: Vec<f32> = if entry.dtype == TensorDtype::Q8_2f {
213                    let col_off = n + rows * 2;
214                    (0..cols)
215                        .map(|i| {
216                            f16_to_f32(u16::from_le_bytes([
217                                bytes[col_off + i * 2],
218                                bytes[col_off + i * 2 + 1],
219                            ]))
220                        })
221                        .collect()
222                } else {
223                    Vec::new()
224                };
225                Ok(Self::Mapped {
226                    model: model.clone(),
227                    idx,
228                    dtype: entry.dtype,
229                    rows,
230                    cols,
231                    row_scale,
232                    col_field,
233                    vbit_offsets: Vec::new(),
234                    repack: q8_repack(bytes, rows, cols),
235                })
236            }
237            // vbit: fused kernel unpacks variable-bit rows from mmap.
238            TensorDtype::Vbit if cols % GROUP_SIZE == 0 => Ok(Self::Mapped {
239                model: model.clone(),
240                idx,
241                dtype: entry.dtype,
242                rows,
243                cols,
244                row_scale: Vec::new(),
245                col_field: Vec::new(),
246                vbit_offsets: vbit_row_offsets(bytes, rows, cols),
247                repack: Vec::new(),
248            }),
249            // vbit_ro (§4.2): the offset table comes straight from the
250            // file — no load-time prefix scan; kernels are shared with
251            // legacy vbit (they consume absolute offsets either way).
252            TensorDtype::VbitRo if cols % GROUP_SIZE == 0 => {
253                let (_, off_off, packed_off) = cortiq_core::quant::vbit_ro_sections(rows, cols);
254                let offsets: Vec<usize> = (0..=rows)
255                    .map(|r| packed_off + cortiq_core::quant::vbit_ro_offset(bytes, off_off, r))
256                    .collect();
257                Ok(Self::Mapped {
258                    model: model.clone(),
259                    idx,
260                    dtype: entry.dtype,
261                    rows,
262                    cols,
263                    row_scale: Vec::new(),
264                    col_field: Vec::new(),
265                    vbit_offsets: offsets,
266                    repack: Vec::new(),
267                })
268            }
269            // q4_block: fused kernel reads nibbles straight from mmap —
270            // a 14B q4 file no longer explodes into ×8 f32 RAM.
271            // q4_tiled (§4.3): interleaved [scale][nibbles] tiles — one
272            // sequential memory stream (measured ×1.66 ARM / ×1.13 AVX2
273            // at kernel level over the split layout).
274            TensorDtype::Q4Tiled if cols % GROUP_SIZE == 0 => Ok(Self::Mapped {
275                model: model.clone(),
276                idx,
277                dtype: entry.dtype,
278                rows,
279                cols,
280                row_scale: Vec::new(),
281                col_field: Vec::new(),
282                vbit_offsets: Vec::new(),
283                repack: Vec::new(),
284            }),
285            // q4tp (§4.10): nibbles from mmap, scale from the row ladder —
286            // 7.3% less file than q4t at the same 4-bit grid.
287            TensorDtype::Q4TiledP if cols % GROUP_SIZE == 0 => Ok(Self::Mapped {
288                model: model.clone(),
289                idx,
290                dtype: entry.dtype,
291                rows,
292                cols,
293                row_scale: Vec::new(),
294                col_field: Vec::new(),
295                vbit_offsets: Vec::new(),
296                repack: Vec::new(),
297            }),
298            // q2tp: 2-bit chunks from mmap, scale from the same row ladder.
299            TensorDtype::Q2TiledP 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            TensorDtype::Q4Block if cols % GROUP_SIZE == 0 => Ok(Self::Mapped {
311                model: model.clone(),
312                idx,
313                dtype: entry.dtype,
314                rows,
315                cols,
316                row_scale: Vec::new(),
317                col_field: Vec::new(),
318                vbit_offsets: Vec::new(),
319                repack: Vec::new(),
320            }),
321            // q1: binary sign-bit tiles from mmap (1-bit-trained models).
322            TensorDtype::Q1 if cols % GROUP_SIZE == 0 => Ok(Self::Mapped {
323                model: model.clone(),
324                idx,
325                dtype: entry.dtype,
326                rows,
327                cols,
328                row_scale: Vec::new(),
329                col_field: Vec::new(),
330                vbit_offsets: Vec::new(),
331                repack: Vec::new(),
332            }),
333            // q1t (ternary + outlier overlay): fused per-row dequant kernel
334            // reads straight from mmap — a 12B q1t stays ~its file size in
335            // RAM instead of dequantizing to ~48 GB of f32.
336            TensorDtype::Q1T if cols % GROUP_SIZE == 0 => Ok(Self::Mapped {
337                model: model.clone(),
338                idx,
339                dtype: entry.dtype,
340                rows,
341                cols,
342                row_scale: Vec::new(),
343                col_field: Vec::new(),
344                vbit_offsets: Vec::new(),
345                repack: Vec::new(),
346            }),
347            // No fused kernel yet → dequantize once (correct, more RAM).
348            _ => {
349                let mut data = vec![0.0f32; rows * cols];
350                cortiq_core::quant::dequant_tensor(entry, bytes, &mut data)?;
351                Ok(Self::from_f32(data, rows, cols))
352            }
353        }
354    }
355
356    /// q1-mapped tensor? (GPU gates: the q1 CPU kernel is
357    /// compute-bound, so offload pays at much smaller shapes than q8.)
358    pub(crate) fn is_q1(&self) -> bool {
359        matches!(
360            self,
361            Self::Mapped {
362                dtype: TensorDtype::Q1,
363                ..
364            }
365        )
366    }
367
368    /// Owned-f32 view (data, rows, cols) — the GDN a/b gate projections
369    /// arrive dequantized (force-f16 in the converter → F32 in RAM).
370    pub(crate) fn f32_parts(&self) -> Option<(&[f32], usize, usize)> {
371        match self {
372            Self::F32 { data, rows, cols } => Some((data, *rows, *cols)),
373            _ => None,
374        }
375    }
376
377    /// (directory idx, rows, cols) of a q1-mapped tensor — the
378    /// whole-block GPU path resolves offsets itself.
379    /// (idx, rows, cols) of a mapped tensor the whole-token GPU graph can drive
380    /// — Q1, Q1T or Q4-block (it resolves the offset and picks the kernel by
381    /// dtype). Q4-block lets a precise down_proj/lm_head stay on-device.
382    /// Named `q1_parts` for historical reasons.
383    pub(crate) fn q1_parts(&self) -> Option<(usize, usize, usize)> {
384        if self.has_prism_contract() {
385            return None;
386        }
387        match self {
388            #[cfg(target_os = "macos")]
389            Self::Mapped {
390                dtype: TensorDtype::Q1T,
391                ..
392            } if !crate::gpu::metal_q1t_enabled() => None,
393            Self::Mapped {
394                idx,
395                dtype:
396                    TensorDtype::Q1
397                    | TensorDtype::Q1T
398                    | TensorDtype::Q4Block
399                    | TensorDtype::Q4Tiled
400                    // Q2TiledP deliberately absent: the Metal graph has no
401                    // q2tp kernel, and advertising it here made the block
402                    // plan truncate mid-run at the first q2tp layer.
403                    | TensorDtype::Q4TiledP
404                    | TensorDtype::Q8Row
405                    | TensorDtype::Q8_2f,
406                rows,
407                cols,
408                ..
409            } => Some((*idx, *rows, *cols)),
410            _ => None,
411        }
412    }
413
414    /// `(directory idx, rows, cols)` for the native Metal token graph.  The
415    /// historical q1 graph gate intentionally refuses every Prism tensor so
416    /// an untransformed q2 payload cannot slip into the resident path.  The
417    /// Metal2 graph is descriptor-aware and admits only the production
418    /// q2tp-affine forward targets; ordinary q1/q4 callers retain the old
419    /// `q1_parts` behaviour.
420    #[cfg(target_os = "macos")]
421    pub(crate) fn metal_graph_parts(&self) -> Option<(usize, usize, usize)> {
422        if let Some((model, idx, kind, _)) = self.graph_weight_descriptor() {
423            let name = &model.tensors[idx].name;
424            let forward = kind == 9 && crate::prism::is_forward_weight(model, name);
425            let affine = kind == 9 && crate::prism::is_affine_target(model, name);
426            if forward && affine {
427                let e = model.tensors.get(idx)?;
428                return Some((idx, *e.shape.first()?, *e.shape.get(1)?));
429            }
430        }
431        self.q1_parts()
432    }
433
434    /// (directory idx, rows, cols) of a q4_tiled mapped tensor. The
435    /// chunk-prefill graph takes it in the same 4-tuple slot as
436    /// `q8_row_parts` with an EMPTY row_scale — q4t carries its scales
437    /// inside the 18-byte tiles, and the empty slice is what tells the
438    /// encoder to reach for the q4t kernels.
439    pub(crate) fn q4t_parts(&self) -> Option<(usize, usize, usize)> {
440        if self.has_prism_contract() {
441            return None;
442        }
443        match self {
444            Self::Mapped {
445                idx,
446                dtype: TensorDtype::Q4Tiled,
447                rows,
448                cols,
449                ..
450            } => Some((*idx, *rows, *cols)),
451            _ => None,
452        }
453    }
454
455    /// (directory idx, rows, cols) of a q4tp mapped tensor. Same empty-scale
456    /// slot as `q4t_parts` in the chunk graph — the encoder tells the two
457    /// apart by the tensor's dtype, not by the slot.
458    pub(crate) fn q4tp_parts(&self) -> Option<(usize, usize, usize)> {
459        if self.has_prism_contract() {
460            return None;
461        }
462        match self {
463            Self::Mapped {
464                idx,
465                dtype: TensorDtype::Q4TiledP,
466                rows,
467                cols,
468                ..
469            } => Some((*idx, *rows, *cols)),
470            _ => None,
471        }
472    }
473
474    /// (directory idx, rows, cols, row_scale) of a plain q8_row mapped
475    /// tensor — the chunk-prefill GPU graph resolves offsets itself.
476    /// q8_2f is excluded on purpose: its column field would need a
477    /// prescale stage on the device.
478    pub(crate) fn q8_row_parts(&self) -> Option<(usize, usize, usize, &[f32])> {
479        if self.has_prism_contract() {
480            return None;
481        }
482        match self {
483            Self::Mapped {
484                idx,
485                dtype: TensorDtype::Q8Row,
486                rows,
487                cols,
488                row_scale,
489                col_field,
490                ..
491            } if col_field.is_empty() => Some((*idx, *rows, *cols, row_scale)),
492            _ => None,
493        }
494    }
495
496    /// The layout this tensor is stored in, when it is mapped from a model.
497    /// The frames branch on it — a q2tp gate against a q4tp down is a real
498    /// combination in the 2-bit profile and needs a different kernel.
499    pub fn model_dtype(&self) -> Option<cortiq_core::TensorDtype> {
500        match self {
501            Self::Mapped { dtype, .. } => Some(*dtype),
502            _ => None,
503        }
504    }
505
506    /// The tensor's index in the model directory, when it is mapped from one.
507    /// The GPU frames bind by index rather than by name — a name lookup per
508    /// layer per token is not free, and the index is what the device cache is
509    /// keyed on anyway.
510    pub fn model_idx(&self) -> Option<usize> {
511        match self {
512            Self::Mapped { idx, .. } => Some(*idx),
513            _ => None,
514        }
515    }
516
517    /// The model this tensor is mapped from, when it is mapped at all. The
518    /// GPU frames need the container to reach the bytes; a QTensor already
519    /// holds it, and threading a second handle down every call site to say
520    /// the same thing invites the two to disagree.
521    pub fn model_arc(&self) -> Option<std::sync::Arc<cortiq_core::CmfModel>> {
522        match self {
523            Self::Mapped { model, .. } => Some(model.clone()),
524            _ => None,
525        }
526    }
527
528    /// Whether this mapped tensor belongs to the Prism/Bonsai transform
529    /// contract.  Device graphs do not carry the descriptor, so callers use
530    /// this conservative predicate to stay on the descriptor-aware CPU path
531    /// instead of silently executing an unrotated matrix.
532    pub(crate) fn has_prism_contract(&self) -> bool {
533        matches!(self, Self::Mapped { model, .. } if crate::prism::has_contract(model))
534    }
535
536    pub fn rows(&self) -> usize {
537        match self {
538            Self::F32 { rows, .. } | Self::Mapped { rows, .. } => *rows,
539        }
540    }
541
542    /// Mapped q4t handle (model + directory index) — the fused GPU FFN
543    /// needs the raw file coordinates of its three projections.
544    pub(crate) fn mapped_q4t(&self) -> Option<(&Arc<CmfModel>, usize)> {
545        if self.has_prism_contract() {
546            return None;
547        }
548        match self {
549            Self::Mapped {
550                model,
551                idx,
552                dtype: TensorDtype::Q4Tiled,
553                ..
554            } => Some((model, *idx)),
555            _ => None,
556        }
557    }
558
559    /// Same slot as `mapped_q4t` for a q4tp tensor — the fused DiT FFN picks
560    /// its kernels by which of the two answers.
561    pub fn mapped_q4tp(&self) -> Option<(&Arc<CmfModel>, usize)> {
562        if self.has_prism_contract() {
563            return None;
564        }
565        match self {
566            Self::Mapped {
567                model,
568                idx,
569                dtype: TensorDtype::Q4TiledP,
570                ..
571            } => Some((model, *idx)),
572            _ => None,
573        }
574    }
575
576    /// (model, tensor idx) for a mapped weight in ANY codec the fused device
577    /// paths can run — four-bit tiled or either int8 layout.
578    ///
579    /// The fused DiT chains asked for `mapped_q4tp` by name, so an eight-bit
580    /// container never reached them and rendered through per-op GEMMs even
581    /// after those kernels learned its codec. The gate is what the codec has
582    /// a device GEMM for, not which codec it is.
583    pub fn mapped_device_gemm(&self) -> Option<(&Arc<CmfModel>, usize)> {
584        if self.has_prism_contract() {
585            return None;
586        }
587        match self {
588            Self::Mapped {
589                model,
590                idx,
591                dtype: TensorDtype::Q4TiledP | TensorDtype::Q8Row | TensorDtype::Q8_2f,
592                ..
593            } => Some((model, *idx)),
594            _ => None,
595        }
596    }
597
598    /// (model, tensor idx) for a q2tp mapped weight — the 2-bit twin of
599    /// `mapped_q4tp`, used by the mixed MoE profile.
600    pub fn mapped_q2tp(&self) -> Option<(&Arc<CmfModel>, usize)> {
601        if self.has_prism_contract() {
602            return None;
603        }
604        match self {
605            Self::Mapped {
606                model,
607                idx,
608                dtype: TensorDtype::Q2TiledP,
609                ..
610            } => Some((model, *idx)),
611            _ => None,
612        }
613    }
614
615    pub fn cols(&self) -> usize {
616        match self {
617            Self::F32 { cols, .. } | Self::Mapped { cols, .. } => *cols,
618        }
619    }
620
621    /// (model, tensor idx) for a q1 mapped weight — the wgpu token graph
622    /// keys its resident VRAM cache by idx. None for any other dtype/kind.
623    pub fn mapped_q1(&self) -> Option<(&std::sync::Arc<CmfModel>, usize)> {
624        if self.has_prism_contract() {
625            return None;
626        }
627        match self {
628            Self::Mapped {
629                model,
630                idx,
631                dtype: TensorDtype::Q1,
632                ..
633            } => Some((model, *idx)),
634            _ => None,
635        }
636    }
637
638    /// (model, idx, kind, row_scale) for a graph-capable mapped weight.
639    /// kind: 0=q8_row (per-row scales), 1=q1, 2=q4_block, 3=q1t
640    /// (tile-embedded, no rs), 5=q4_tiled, 6=q4tp, 7=q8_2f (both scale
641    /// planes live inside the tensor). None only for `vbit`.
642    ///
643    /// The old comment here claimed q4_block was unhandled while the arm
644    /// right below mapped it, and it named q8_2f as unhandled after that
645    /// stopped being true — a stale comment on this function is how a
646    /// model silently loses the graph, so it is worth keeping honest.
647    pub fn graph_weight(&self) -> Option<(&std::sync::Arc<CmfModel>, usize, u8, &[f32])> {
648        if self.has_prism_contract() {
649            return None;
650        }
651        self.graph_weight_descriptor()
652    }
653
654    /// Descriptor-aware graph handle used only by the Prism token graph.
655    /// Ordinary graph callers continue to use [`graph_weight`] and therefore
656    /// remain fail-closed until they provide the same explicit transform
657    /// contract.
658    pub(crate) fn graph_weight_descriptor(
659        &self,
660    ) -> Option<(&std::sync::Arc<CmfModel>, usize, u8, &[f32])> {
661        match self {
662            Self::Mapped {
663                model,
664                idx,
665                dtype: TensorDtype::Q8Row,
666                row_scale,
667                ..
668            } => Some((model, *idx, 0, row_scale.as_slice())),
669            Self::Mapped {
670                model,
671                idx,
672                dtype: TensorDtype::Q1,
673                ..
674            } => Some((model, *idx, 1, &[])),
675            // Q4Tiled is kind 5, NOT 2: both carried 2 historically, and
676            // the wgpu token graph fed 18B interleaved tiles to the
677            // split-layout q4b kernel — garbage output on q4t models
678            // (caught by an end-to-end answer check on real Vulkan).
679            Self::Mapped {
680                model,
681                idx,
682                dtype: TensorDtype::Q4Tiled,
683                ..
684            } => Some((model, *idx, 5, &[])),
685            // Kind 6, not 5: q4tp's nibble stride and scale planes differ,
686            // and feeding them to the q4t kernel is exactly the mistake that
687            // produced garbage when Q4Tiled shared kind 2 with Q4Block.
688            Self::Mapped {
689                model,
690                idx,
691                dtype: TensorDtype::Q4TiledP,
692                ..
693            } => Some((model, *idx, 6, &[])),
694            Self::Mapped {
695                model,
696                idx,
697                dtype: TensorDtype::Q4Block,
698                ..
699            } => Some((model, *idx, 2, &[])),
700            // q8_2f carries BOTH scale planes after the int8 body (rows
701            // f16, then cols f16), so the graph takes the whole tensor
702            // and the kernel reads them where they lie — no host-side
703            // prescale, which is what the per-op path does instead.
704            Self::Mapped {
705                model,
706                idx,
707                dtype: TensorDtype::Q8_2f,
708                ..
709            } => Some((model, *idx, 7, &[])),
710            Self::Mapped {
711                model,
712                idx,
713                dtype: TensorDtype::Q1T,
714                ..
715            } => Some((model, *idx, 3, &[])),
716            // Kind 9: the 2-bit plane on the q4tp ladder (dense FFN gate/up
717            // of the q2tp profile). Its own kernel — 8 bytes a group where
718            // q4tp has 16, and rung 0 is the exact zero.
719            Self::Mapped {
720                model,
721                idx,
722                dtype: TensorDtype::Q2TiledP,
723                ..
724            } => Some((model, *idx, 9, &[])),
725            _ => None,
726        }
727    }
728
729    /// Dense f32 view — only for owned tensors. Masked/sparse execution
730    /// paths require it; quantized weights don't support masks yet.
731    pub fn as_f32(&self) -> Option<&[f32]> {
732        match self {
733            Self::F32 { data, .. } => Some(data),
734            Self::Mapped { .. } => None,
735        }
736    }
737
738    fn quant_bytes(&self) -> &[u8] {
739        match self {
740            Self::Mapped { model, idx, .. } => model.entry_bytes(&model.tensors[*idx]),
741            Self::F32 { .. } => unreachable!("quant_bytes on F32"),
742        }
743    }
744
745    /// Dequantize one row into `dst` (embedding lookup).
746    pub fn row_f32(&self, r: usize, dst: &mut [f32]) {
747        let cols = self.cols();
748        debug_assert_eq!(dst.len(), cols);
749        match self {
750            Self::F32 { data, .. } => dst.copy_from_slice(&data[r * cols..(r + 1) * cols]),
751            Self::Mapped {
752                model,
753                idx,
754                dtype,
755                row_scale,
756                col_field,
757                vbit_offsets,
758                ..
759            } => {
760                if *dtype == TensorDtype::Q4Tiled {
761                    let bytes = self.quant_bytes();
762                    let gpr = cols / GROUP_SIZE;
763                    for gi in 0..gpr {
764                        let tile = &bytes[(r * gpr + gi) * Q4_TILE..(r * gpr + gi + 1) * Q4_TILE];
765                        let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
766                        for (k, &b) in tile[2..].iter().enumerate() {
767                            dst[gi * GROUP_SIZE + k * 2] = ((b & 0x0F) as f32 - 8.0) * s;
768                            dst[gi * GROUP_SIZE + k * 2 + 1] = (((b >> 4) & 0x0F) as f32 - 8.0) * s;
769                        }
770                    }
771                    if crate::prism::is_inverse_embedding(model, &model.tensors[*idx].name) {
772                        crate::prism::inverse_embedding(model, dst);
773                    }
774                    return;
775                }
776                if *dtype == TensorDtype::Q4TiledP {
777                    let bytes = self.quant_bytes();
778                    let gpr = cols / GROUP_SIZE;
779                    let v = Q4tpView::new(bytes, self.rows(), cols);
780                    let mut sc = vec![0f32; gpr];
781                    v.scales_into(r, gpr, &mut sc);
782                    for gi in 0..gpr {
783                        let tile = &v.nib[(r * gpr + gi) * Q4TP_NIB..(r * gpr + gi + 1) * Q4TP_NIB];
784                        let s = sc[gi];
785                        for (k, &b) in tile.iter().enumerate() {
786                            dst[gi * GROUP_SIZE + k * 2] = ((b & 0x0F) as f32 - 8.0) * s;
787                            dst[gi * GROUP_SIZE + k * 2 + 1] = (((b >> 4) & 0x0F) as f32 - 8.0) * s;
788                        }
789                    }
790                    if crate::prism::is_inverse_embedding(model, &model.tensors[*idx].name) {
791                        crate::prism::inverse_embedding(model, dst);
792                    }
793                    return;
794                }
795                if *dtype == TensorDtype::Q2TiledP {
796                    let bytes = self.quant_bytes();
797                    let gpr = cols / GROUP_SIZE;
798                    let v = Q4tpView::new_q2(bytes, self.rows(), cols);
799                    let mut sc = vec![0f32; gpr];
800                    v.scales_into(r, gpr, &mut sc);
801                    for gi in 0..gpr {
802                        let ch =
803                            &v.nib[(r * gpr + gi) * Q2TP_CHUNK..(r * gpr + gi + 1) * Q2TP_CHUNK];
804                        let s = sc[gi];
805                        for (k, &b) in ch.iter().enumerate() {
806                            for j in 0..4 {
807                                let center = if crate::prism::is_affine_target(
808                                    model,
809                                    &model.tensors[*idx].name,
810                                ) {
811                                    1.0
812                                } else {
813                                    1.5
814                                };
815                                dst[gi * GROUP_SIZE + k * 4 + j] =
816                                    (((b >> (2 * j)) & 3) as f32 - center) * s;
817                            }
818                        }
819                    }
820                    if crate::prism::is_inverse_embedding(model, &model.tensors[*idx].name) {
821                        crate::prism::inverse_embedding(model, dst);
822                    }
823                    return;
824                }
825                if *dtype == TensorDtype::Q4Block {
826                    let (packed, scales) = q4_split(self.quant_bytes(), self.rows(), cols);
827                    let gpr = cols / GROUP_SIZE;
828                    for gi in 0..gpr {
829                        let g = r * gpr + gi;
830                        let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
831                        for (k, &b) in packed[g * 16..(g + 1) * 16].iter().enumerate() {
832                            dst[gi * GROUP_SIZE + k * 2] = ((b & 0x0F) as f32 - 8.0) * s;
833                            dst[gi * GROUP_SIZE + k * 2 + 1] = (((b >> 4) & 0x0F) as f32 - 8.0) * s;
834                        }
835                    }
836                    if crate::prism::is_inverse_embedding(model, &model.tensors[*idx].name) {
837                        crate::prism::inverse_embedding(model, dst);
838                    }
839                    return;
840                }
841                if *dtype == TensorDtype::Q1 {
842                    let bytes = self.quant_bytes();
843                    let gpr = cols / GROUP_SIZE;
844                    for gi in 0..gpr {
845                        let tile = &bytes[(r * gpr + gi) * Q1_TILE..(r * gpr + gi + 1) * Q1_TILE];
846                        let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
847                        for (j, &b) in tile[2..].iter().enumerate() {
848                            for k in 0..8 {
849                                dst[gi * GROUP_SIZE + j * 8 + k] =
850                                    (((b >> k) & 1) as f32 * 2.0 - 1.0) * s;
851                            }
852                        }
853                    }
854                    if crate::prism::is_inverse_embedding(model, &model.tensors[*idx].name) {
855                        crate::prism::inverse_embedding(model, dst);
856                    }
857                    return;
858                }
859                if *dtype == TensorDtype::Q1T {
860                    let bytes = self.quant_bytes();
861                    let gpr = cols / GROUP_SIZE;
862                    let base_len = self.rows() * gpr * cortiq_core::quant::Q1T_TILE;
863                    for gi in 0..gpr {
864                        let off = (r * gpr + gi) * cortiq_core::quant::Q1T_TILE;
865                        let s = cortiq_core::quant::f16_to_f32(u16::from_le_bytes([
866                            bytes[off],
867                            bytes[off + 1],
868                        ]));
869                        let codes = &bytes[off + 2..off + cortiq_core::quant::Q1T_TILE];
870                        for k in 0..GROUP_SIZE {
871                            dst[gi * GROUP_SIZE + k] = match cortiq_core::quant::q1t_code(codes, k)
872                            {
873                                1 => s,
874                                2 => -s,
875                                _ => 0.0,
876                            };
877                        }
878                    }
879                    // Overlay
880                    let rows = self.rows();
881                    let entries = base_len + (rows + 1) * 4;
882                    if entries <= bytes.len() {
883                        let ptrs = &bytes[base_len..base_len + (rows + 1) * 4];
884                        let r0 = u32::from_le_bytes([
885                            ptrs[r * 4],
886                            ptrs[r * 4 + 1],
887                            ptrs[r * 4 + 2],
888                            ptrs[r * 4 + 3],
889                        ]) as usize;
890                        let r1 = u32::from_le_bytes([
891                            ptrs[(r + 1) * 4],
892                            ptrs[(r + 1) * 4 + 1],
893                            ptrs[(r + 1) * 4 + 2],
894                            ptrs[(r + 1) * 4 + 3],
895                        ]) as usize;
896                        let off = entries + r0 * 4;
897                        for i in 0..r1 - r0 {
898                            let item = &bytes[off + i * 4..off + i * 4 + 4];
899                            let c = u16::from_le_bytes([item[0], item[1]]) as usize;
900                            let v = cortiq_core::quant::f16_to_f32(u16::from_le_bytes([
901                                item[2], item[3],
902                            ]));
903                            if c < cols {
904                                dst[c] = v;
905                            }
906                        }
907                    }
908                    if crate::prism::is_inverse_embedding(model, &model.tensors[*idx].name) {
909                        crate::prism::inverse_embedding(model, dst);
910                    }
911                    return;
912                }
913                if matches!(dtype, TensorDtype::Vbit | TensorDtype::VbitRo) {
914                    let bytes = self.quant_bytes();
915                    let rows = self.rows();
916                    let ng = cols / GROUP_SIZE;
917                    let bits = &bytes[..rows];
918                    let sc_off = rows;
919                    // Precomputed at load — embedding lookup used to scan
920                    // the bit-widths of every preceding row (O(token_id)).
921                    let off = vbit_offsets[r];
922                    let b = bits[r] as usize;
923                    let l = ((1usize << (b - 1)) - 1) as f32;
924                    let data = &bytes[off..];
925                    let (mut acc, mut nbits, mut byte_idx) = (0u64, 0usize, 0usize);
926                    for (i, d) in dst.iter_mut().enumerate() {
927                        while nbits < b {
928                            acc = (acc << 8) | data[byte_idx] as u64;
929                            byte_idx += 1;
930                            nbits += 8;
931                        }
932                        let u = ((acc >> (nbits - b)) & ((1u64 << b) - 1)) as f32;
933                        nbits -= b;
934                        let so = (r * ng + i / GROUP_SIZE) * 2;
935                        let sv = f16_to_f32(u16::from_le_bytes([
936                            bytes[sc_off + so],
937                            bytes[sc_off + so + 1],
938                        ]));
939                        *d = (u - l) * sv;
940                    }
941                    if crate::prism::is_inverse_embedding(model, &model.tensors[*idx].name) {
942                        crate::prism::inverse_embedding(model, dst);
943                    }
944                    return;
945                }
946                let q = &self.quant_bytes()[r * cols..(r + 1) * cols];
947                let s = row_scale[r];
948                match dtype {
949                    TensorDtype::Q8Row => {
950                        for (d, &b) in dst.iter_mut().zip(q) {
951                            *d = (b as i8) as f32 * s;
952                        }
953                    }
954                    TensorDtype::Q8_2f => {
955                        for (i, (d, &b)) in dst.iter_mut().zip(q).enumerate() {
956                            *d = (b as i8) as f32 * s * col_field[i];
957                        }
958                    }
959                    _ => unreachable!(),
960                }
961                if crate::prism::is_inverse_embedding(model, &model.tensors[*idx].name) {
962                    crate::prism::inverse_embedding(model, dst);
963                }
964            }
965        }
966    }
967
968    /// Can this tensor's columns be read cheaply (for sparse down_proj)?
969    /// True for F32/Q8Row/Q8_2f (per-row scale, direct strided access);
970    /// false for group-packed q4/vbit (column access would unpack whole
971    /// groups — sparse execution falls back to f32 for those).
972    pub fn sparse_col_ok(&self) -> bool {
973        match self {
974            Self::F32 { .. } => true,
975            Self::Mapped { dtype, .. } => {
976                matches!(dtype, TensorDtype::Q8Row | TensorDtype::Q8_2f)
977            }
978        }
979    }
980
981    /// down_proj [hidden, inter]: accumulate `w · col(c)` into `out`
982    /// [hidden] — reads ONLY column `c` (one neuron) from the mmap,
983    /// no full-matrix dequant. `out[k] += w · down[k, c]`.
984    pub fn add_col_scaled(&self, c: usize, w: f32, out: &mut [f32]) {
985        let inter = self.cols();
986        let hidden = self.rows();
987        debug_assert_eq!(out.len(), hidden);
988        match self {
989            Self::F32 { data, .. } => {
990                for (k, o) in out.iter_mut().enumerate() {
991                    *o += w * data[k * inter + c];
992                }
993            }
994            Self::Mapped {
995                dtype,
996                row_scale,
997                col_field,
998                ..
999            } => {
1000                let q = self.quant_bytes();
1001                let colf = if *dtype == TensorDtype::Q8_2f {
1002                    col_field[c]
1003                } else {
1004                    1.0
1005                };
1006                let wc = w * colf;
1007                for (k, o) in out.iter_mut().enumerate() {
1008                    let b = q[k * inter + c] as i8 as f32;
1009                    *o += wc * b * row_scale[k];
1010                }
1011            }
1012        }
1013    }
1014
1015    /// Touch the head of row `r` so the DRAM latency of the next
1016    /// neuron's weights overlaps the current one's arithmetic.
1017    ///
1018    /// Scattered rows are what per-token sparsity reads, and a 2 KB
1019    /// stride is past what the hardware prefetcher follows: without this
1020    /// every row starts with a cold miss that nothing hides. One touch
1021    /// per 512 bytes is enough — the rest of the row is a sequential run
1022    /// the prefetcher does pick up.
1023    #[inline]
1024    pub fn prefetch_row(&self, r: usize) {
1025        let Self::Mapped { dtype, .. } = self else {
1026            return;
1027        };
1028        if !matches!(dtype, TensorDtype::Q8Row | TensorDtype::Q8_2f) {
1029            return;
1030        }
1031        let cols = self.cols();
1032        let q = self.quant_bytes();
1033        let (a, b) = (r * cols, (r + 1) * cols);
1034        if b > q.len() {
1035            return;
1036        }
1037        let mut j = a;
1038        while j < b {
1039            unsafe { std::ptr::read_volatile(q.as_ptr().add(j)) };
1040            j += 512;
1041        }
1042    }
1043
1044    /// `out += w · row(r)` — the transposed twin of `add_col_scaled`.
1045    ///
1046    /// A neuron's `down` weights are a COLUMN of `[hidden, inter]`, and a
1047    /// column is strided: reading one costs a cache line per element, so
1048    /// per-neuron dynamic sparsity saves arithmetic and no bytes. Stored
1049    /// transposed (`down_proj.t.weight`, `[inter, hidden]`) the same
1050    /// weights are a contiguous ROW, and this accumulate reads exactly
1051    /// the neurons the token asked for.
1052    pub fn add_row_scaled(&self, r: usize, w: f32, out: &mut [f32], scratch: &mut [f32]) {
1053        let cols = self.cols();
1054        debug_assert_eq!(out.len(), cols);
1055        match self {
1056            Self::F32 { data, .. } => {
1057                let row = &data[r * cols..(r + 1) * cols];
1058                for (o, v) in out.iter_mut().zip(row) {
1059                    *o += w * v;
1060                }
1061            }
1062            Self::Mapped {
1063                dtype,
1064                row_scale,
1065                col_field,
1066                ..
1067            } => match dtype {
1068                TensorDtype::Q8Row => {
1069                    let q = &self.quant_bytes()[r * cols..(r + 1) * cols];
1070                    let ws = w * row_scale[r];
1071                    let row: &[i8] =
1072                        unsafe { std::slice::from_raw_parts(q.as_ptr() as *const i8, q.len()) };
1073                    axpy_i8_f32(out, row, ws);
1074                }
1075                TensorDtype::Q8_2f => {
1076                    let q = &self.quant_bytes()[r * cols..(r + 1) * cols];
1077                    let ws = w * row_scale[r];
1078                    for ((o, b), c) in out.iter_mut().zip(q).zip(col_field) {
1079                        *o += ws * c * (*b as i8 as f32);
1080                    }
1081                }
1082                _ => {
1083                    self.row_f32(r, scratch);
1084                    for (o, v) in out.iter_mut().zip(scratch.iter()) {
1085                        *o += w * v;
1086                    }
1087                }
1088            },
1089        }
1090    }
1091
1092    /// Dot of row `r` with `x` (gate/up active-neuron path). Reads only
1093    /// row `r` from the mmap — no full dequant. q4/vbit dequant the row
1094    /// into `scratch` first (rare for active-FFN weights).
1095    pub fn row_dot(&self, r: usize, x: &[f32], scratch: &mut [f32]) -> f32 {
1096        let cols = self.cols();
1097        match self {
1098            Self::F32 { data, .. } => {
1099                let row = &data[r * cols..(r + 1) * cols];
1100                row.iter().zip(x).map(|(w, v)| w * v).sum()
1101            }
1102            Self::Mapped {
1103                model,
1104                idx,
1105                dtype,
1106                row_scale,
1107                col_field,
1108                ..
1109            } => {
1110                let prism_forward =
1111                    crate::prism::is_forward_weight(model, &model.tensors[*idx].name);
1112                if prism_forward {
1113                    let transformed = crate::prism::forward(model, &x[..cols]);
1114                    let gpr = cols / GROUP_SIZE;
1115                    match dtype {
1116                        TensorDtype::Q2TiledP => {
1117                            let v = Q4tpView::new_q2(self.quant_bytes(), self.rows(), cols);
1118                            let mut sc = vec![0f32; gpr];
1119                            v.scales_into(r, gpr, &mut sc);
1120                            if crate::prism::is_affine_target(model, &model.tensors[*idx].name) {
1121                                return q2tp_affine_row_exact(v.nib, r, gpr, &transformed, &sc);
1122                            }
1123                            return q2tp_row_exact(v.nib, r, gpr, &transformed, &sc);
1124                        }
1125                        _ => {
1126                            self.row_f32(r, scratch);
1127                            return scratch.iter().zip(&transformed).map(|(w, v)| w * v).sum();
1128                        }
1129                    }
1130                }
1131                match dtype {
1132                    TensorDtype::Q8Row => {
1133                        let q = &self.quant_bytes()[r * cols..(r + 1) * cols];
1134                        dot_i8_f32(q, x) * row_scale[r]
1135                    }
1136                    TensorDtype::Q8_2f => {
1137                        let q = &self.quant_bytes()[r * cols..(r + 1) * cols];
1138                        dot_i8_col_f32(q, x, col_field) * row_scale[r]
1139                    }
1140                    _ => {
1141                        self.row_f32(r, scratch);
1142                        scratch.iter().zip(x).map(|(w, v)| w * v).sum()
1143                    }
1144                }
1145            }
1146        }
1147    }
1148
1149    /// `out = W · x` (row-major). F32 delegates to the historical
1150    /// bit-exact path; Mapped runs the fused int8 kernel.
1151    pub fn matvec(&self, x: &[f32], out: &mut [f32], pool: Option<&Pool>) {
1152        match self {
1153            // NOTE: `out.len()` DRIVES this arm — it computes that many rows,
1154            // and `x.len()` is the stride. A short `out` is legitimate here,
1155            // which is why the check below lives in the Mapped arm only.
1156            Self::F32 { data, .. } => matvec_rows(pool, data, x, out),
1157            Self::Mapped {
1158                model,
1159                idx,
1160                dtype,
1161                rows,
1162                cols,
1163                row_scale,
1164                col_field,
1165                vbit_offsets,
1166                repack,
1167            } => {
1168                let _ = (model, idx);
1169                // Every kernel below writes `rows` entries through a raw
1170                // pointer, so a short `out` is an out-of-bounds WRITE, not a
1171                // wrong answer: it scribbles on the allocator's metadata and
1172                // the process aborts much later, somewhere innocent
1173                // (`double free or corruption`, `corrupted double-linked
1174                // list`). The debug_assert two of the kernels carried is
1175                // compiled out of the release — exactly the build where it
1176                // matters. Fail here instead, while the caller is still on
1177                // the stack to be named.
1178                assert!(
1179                    out.len() >= *rows && x.len() >= *cols,
1180                    "matvec {rows}x{cols}: out {} (need {rows}), x {} (need {cols})",
1181                    out.len(),
1182                    x.len(),
1183                );
1184                let prism_forward =
1185                    crate::prism::is_forward_weight(model, &model.tensors[*idx].name);
1186                if *dtype == TensorDtype::Q2TiledP
1187                    && std::env::var("CMF_Q2TP_TRACE").as_deref() == Ok("1")
1188                {
1189                    use std::sync::atomic::{AtomicUsize, Ordering};
1190                    static N: AtomicUsize = AtomicUsize::new(0);
1191                    let n = N.fetch_add(1, Ordering::Relaxed);
1192                    if n < 128 {
1193                        eprintln!(
1194                            "q2tp-dispatch #{n} name={} prism={} rows={} cols={} gpu={} optin={} layer={}",
1195                            model.tensors[*idx].name,
1196                            prism_forward,
1197                            rows,
1198                            cols,
1199                            crate::gpu::enabled_here(),
1200                            crate::gpu::q2tp_gpu_opt_in(),
1201                            crate::gpu::cur_layer(),
1202                        );
1203                    }
1204                }
1205                // Prism stores every manifest-listed forward matrix in the
1206                // signed-Hadamard basis.  The q2tp WGSL path receives that
1207                // transformed vector and an explicit affine bit; codecs
1208                // without a descriptor-aware kernel remain on CPU below.
1209                if prism_forward {
1210                    let transformed = crate::prism::forward(model, &x[..*cols]);
1211                    match dtype {
1212                        TensorDtype::Q4Block => {
1213                            q4matvec(self.quant_bytes(), &transformed, *rows, *cols, out, pool)
1214                        }
1215                        TensorDtype::Q4Tiled => {
1216                            q4t_matvec(self.quant_bytes(), &transformed, *rows, *cols, out, pool)
1217                        }
1218                        TensorDtype::Q4TiledP => {
1219                            q4tp_matvec(self.quant_bytes(), &transformed, *rows, *cols, out, pool)
1220                        }
1221                        TensorDtype::Q2TiledP => {
1222                            let affine =
1223                                crate::prism::is_affine_target(model, &model.tensors[*idx].name);
1224                            if *rows * *cols >= 8_388_608
1225                                && crate::gpu::enabled_here()
1226                                && crate::gpu::q2tp_gpu_opt_in()
1227                            {
1228                                let gpu_ok = if affine {
1229                                    crate::gpu::q2tp_affine_matvec(
1230                                        model,
1231                                        *idx,
1232                                        &transformed,
1233                                        *rows,
1234                                        *cols,
1235                                        out,
1236                                    )
1237                                } else {
1238                                    crate::gpu::q2tp_matvec(
1239                                        model,
1240                                        *idx,
1241                                        &transformed,
1242                                        *rows,
1243                                        *cols,
1244                                        out,
1245                                    )
1246                                };
1247                                if gpu_ok {
1248                                    return;
1249                                }
1250                            }
1251                            if affine {
1252                                q2tp_affine_matvec(
1253                                    self.quant_bytes(),
1254                                    &transformed,
1255                                    *rows,
1256                                    *cols,
1257                                    out,
1258                                    pool,
1259                                )
1260                            } else {
1261                                q2tp_matvec(
1262                                    self.quant_bytes(),
1263                                    &transformed,
1264                                    *rows,
1265                                    *cols,
1266                                    out,
1267                                    pool,
1268                                )
1269                            }
1270                        }
1271                        TensorDtype::Q1 => {
1272                            q1_matvec(self.quant_bytes(), &transformed, *rows, *cols, out, pool)
1273                        }
1274                        TensorDtype::Q1T => {
1275                            q1t_matvec(self.quant_bytes(), &transformed, *rows, *cols, out, pool)
1276                        }
1277                        TensorDtype::Vbit | TensorDtype::VbitRo => vbitmatvec(
1278                            self.quant_bytes(),
1279                            vbit_offsets,
1280                            &transformed,
1281                            *rows,
1282                            *cols,
1283                            out,
1284                            pool,
1285                        ),
1286                        TensorDtype::Q8Row | TensorDtype::Q8_2f => qmatvec(
1287                            self.quant_bytes(),
1288                            repack,
1289                            row_scale,
1290                            &transformed,
1291                            col_field,
1292                            *dtype,
1293                            *rows,
1294                            *cols,
1295                            out,
1296                            pool,
1297                        ),
1298                        _ => unreachable!("unsupported mapped Prism dtype {dtype:?}"),
1299                    }
1300                    return;
1301                }
1302                if *dtype == TensorDtype::Q4Block {
1303                    // GPU route (wgpu q4b kernel) for large q4_block matvecs —
1304                    // gives NVIDIA/AMD/Intel q4 models a GPU path. Probe keeps
1305                    // the winner; Metal returns false → the CPU kernel below.
1306                    if *rows * *cols >= 8_388_608 && crate::gpu::enabled_here() {
1307                        let t0 = std::time::Instant::now();
1308                        match crate::gpu::probe_arm(crate::gpu::OpClass::Matvec) {
1309                            crate::gpu::ProbeArm::Gpu => {
1310                                if crate::gpu::q4b_matvec(model, *idx, x, *rows, *cols, out) {
1311                                    crate::gpu::probe_record(
1312                                        crate::gpu::OpClass::Matvec,
1313                                        true,
1314                                        t0.elapsed(),
1315                                    );
1316                                    return;
1317                                }
1318                            }
1319                            crate::gpu::ProbeArm::CpuTimed => {
1320                                q4matvec(self.quant_bytes(), x, *rows, *cols, out, pool);
1321                                crate::gpu::probe_record(
1322                                    crate::gpu::OpClass::Matvec,
1323                                    false,
1324                                    t0.elapsed(),
1325                                );
1326                                return;
1327                            }
1328                            crate::gpu::ProbeArm::Cpu => {}
1329                        }
1330                    }
1331                    q4matvec(self.quant_bytes(), x, *rows, *cols, out, pool);
1332                    return;
1333                }
1334                if *dtype == TensorDtype::Q4Tiled {
1335                    // GPU route for large q4t matvecs — the lm_head class,
1336                    // same shape as the q4tp arm below. The probe keeps the
1337                    // winner; a backend without the kernel refuses and the
1338                    // CPU path stays.
1339                    if *rows * *cols >= 8_388_608 && crate::gpu::enabled_here() {
1340                        let t0 = std::time::Instant::now();
1341                        let cls = crate::gpu::matvec_class(*rows, *cols);
1342                        match crate::gpu::probe_arm(cls) {
1343                            crate::gpu::ProbeArm::Gpu => {
1344                                if crate::gpu::q4t_matvec(model, *idx, x, *rows, *cols, out) {
1345                                    crate::gpu::probe_record(cls, true, t0.elapsed());
1346                                    return;
1347                                }
1348                            }
1349                            crate::gpu::ProbeArm::CpuTimed => {
1350                                q4t_matvec(self.quant_bytes(), x, *rows, *cols, out, pool);
1351                                crate::gpu::probe_record(cls, false, t0.elapsed());
1352                                return;
1353                            }
1354                            crate::gpu::ProbeArm::Cpu => {}
1355                        }
1356                    }
1357                    q4t_matvec(self.quant_bytes(), x, *rows, *cols, out, pool);
1358                    return;
1359                }
1360                if *dtype == TensorDtype::Q4TiledP {
1361                    // GPU route for large q4tp matvecs — the lm_head class.
1362                    // On a q4tp checkpoint the head is the biggest single
1363                    // host matvec left in the decode step, and the batched
1364                    // kernel at b=1 already exists on both backends. Probe
1365                    // keeps the winner, same as q4_block above.
1366                    if *rows * *cols >= 8_388_608 && crate::gpu::enabled_here() {
1367                        let t0 = std::time::Instant::now();
1368                        let cls = crate::gpu::matvec_class(*rows, *cols);
1369                        match crate::gpu::probe_arm(cls) {
1370                            crate::gpu::ProbeArm::Gpu => {
1371                                if crate::gpu::q4tp_matvec(model, *idx, x, *rows, *cols, out) {
1372                                    crate::gpu::probe_record(cls, true, t0.elapsed());
1373                                    return;
1374                                }
1375                            }
1376                            crate::gpu::ProbeArm::CpuTimed => {
1377                                q4tp_matvec(self.quant_bytes(), x, *rows, *cols, out, pool);
1378                                crate::gpu::probe_record(cls, false, t0.elapsed());
1379                                return;
1380                            }
1381                            crate::gpu::ProbeArm::Cpu => {}
1382                        }
1383                    }
1384                    q4tp_matvec(self.quant_bytes(), x, *rows, *cols, out, pool);
1385                    return;
1386                }
1387                if *dtype == TensorDtype::Q2TiledP {
1388                    q2tp_matvec(self.quant_bytes(), x, *rows, *cols, out, pool);
1389                    return;
1390                }
1391                if *dtype == TensorDtype::Q1 {
1392                    // GPU route for large q1 matvecs (out_proj / lm_head
1393                    // class): the CPU q1 kernel is load-port-bound at
1394                    // ~4 GB/s/core, the GPU one is bandwidth-bound — the
1395                    // probe measures both arms and keeps the winner.
1396                    if *rows * *cols >= 8_388_608 && crate::gpu::enabled_here() {
1397                        let t0 = std::time::Instant::now();
1398                        let arm = if crate::gpu::q1_force() {
1399                            crate::gpu::ProbeArm::Gpu
1400                        } else {
1401                            crate::gpu::probe_arm(crate::gpu::OpClass::Matvec)
1402                        };
1403                        match arm {
1404                            crate::gpu::ProbeArm::Gpu => {
1405                                if crate::gpu::q1_matvec(model, *idx, x, *rows, *cols, out) {
1406                                    crate::gpu::probe_record(
1407                                        crate::gpu::OpClass::Matvec,
1408                                        true,
1409                                        t0.elapsed(),
1410                                    );
1411                                    return;
1412                                }
1413                            }
1414                            crate::gpu::ProbeArm::CpuTimed => {
1415                                q1_matvec(self.quant_bytes(), x, *rows, *cols, out, pool);
1416                                crate::gpu::probe_record(
1417                                    crate::gpu::OpClass::Matvec,
1418                                    false,
1419                                    t0.elapsed(),
1420                                );
1421                                return;
1422                            }
1423                            crate::gpu::ProbeArm::Cpu => {}
1424                        }
1425                    }
1426                    q1_matvec(self.quant_bytes(), x, *rows, *cols, out, pool);
1427                    return;
1428                }
1429                if *dtype == TensorDtype::Q1T {
1430                    // GPU route for large q1t matvecs: the ternary BASE dot runs
1431                    // on the GPU (load-port-bound on CPU, like q1), then the
1432                    // sparse overlay is added on the CPU. Probe keeps the winner.
1433                    if *rows * *cols >= 8_388_608 && crate::gpu::enabled_here() {
1434                        let t0 = std::time::Instant::now();
1435                        match crate::gpu::probe_arm(crate::gpu::OpClass::Matvec) {
1436                            crate::gpu::ProbeArm::Gpu => {
1437                                if crate::gpu::q1t_matvec(model, *idx, x, *rows, *cols, out) {
1438                                    q1t_add_overlay(self.quant_bytes(), x, *rows, *cols, out, pool);
1439                                    crate::gpu::probe_record(
1440                                        crate::gpu::OpClass::Matvec,
1441                                        true,
1442                                        t0.elapsed(),
1443                                    );
1444                                    return;
1445                                }
1446                            }
1447                            crate::gpu::ProbeArm::CpuTimed => {
1448                                q1t_matvec(self.quant_bytes(), x, *rows, *cols, out, pool);
1449                                crate::gpu::probe_record(
1450                                    crate::gpu::OpClass::Matvec,
1451                                    false,
1452                                    t0.elapsed(),
1453                                );
1454                                return;
1455                            }
1456                            crate::gpu::ProbeArm::Cpu => {}
1457                        }
1458                    }
1459                    q1t_matvec(self.quant_bytes(), x, *rows, *cols, out, pool);
1460                    return;
1461                }
1462                if matches!(dtype, TensorDtype::Vbit | TensorDtype::VbitRo) {
1463                    vbitmatvec(self.quant_bytes(), vbit_offsets, x, *rows, *cols, out, pool);
1464                    return;
1465                }
1466                let xs = prescale(x, col_field, *dtype);
1467                // D5: large q8 matrices (lm_head-class) — hybrid
1468                // CPU∥GPU: split the rows, both sides compute
1469                // SIMULTANEOUSLY (same math, shared prescale).
1470                // GPU share: CMF_GPU_SPLIT (0..1, default 0.5).
1471                if *rows >= crate::gpu::min_rows()
1472                    && matches!(dtype, TensorDtype::Q8Row | TensorDtype::Q8_2f)
1473                    && gpu_lmhead_enabled()
1474                    && crate::gpu::enabled_here()
1475                {
1476                    // Runtime probe: alternate the hybrid against the
1477                    // pure-CPU matvec, keep whichever is faster HERE.
1478                    let t0 = std::time::Instant::now();
1479                    match crate::gpu::probe_arm(crate::gpu::OpClass::Matvec) {
1480                        crate::gpu::ProbeArm::Gpu => {}
1481                        crate::gpu::ProbeArm::CpuTimed => {
1482                            qmatvec(
1483                                self.quant_bytes(),
1484                                repack,
1485                                row_scale,
1486                                x,
1487                                col_field,
1488                                *dtype,
1489                                *rows,
1490                                *cols,
1491                                out,
1492                                pool,
1493                            );
1494                            crate::gpu::probe_record(
1495                                crate::gpu::OpClass::Matvec,
1496                                false,
1497                                t0.elapsed(),
1498                            );
1499                            return;
1500                        }
1501                        crate::gpu::ProbeArm::Cpu => {
1502                            qmatvec(
1503                                self.quant_bytes(),
1504                                repack,
1505                                row_scale,
1506                                x,
1507                                col_field,
1508                                *dtype,
1509                                *rows,
1510                                *cols,
1511                                out,
1512                                pool,
1513                            );
1514                            return;
1515                        }
1516                    }
1517                    let frac = gpu_split_frac();
1518                    let cpu_rows = ((*rows as f32) * (1.0 - frac)) as usize;
1519                    let (out_cpu, out_gpu) = out.split_at_mut(cpu_rows);
1520                    let bytes = self.quant_bytes();
1521                    let ok = std::thread::scope(|sc| {
1522                        let g = sc.spawn(|| {
1523                            crate::gpu::q8_matvec_range(
1524                                model,
1525                                *idx,
1526                                cpu_rows,
1527                                &row_scale[cpu_rows..],
1528                                &xs,
1529                                *rows - cpu_rows,
1530                                *cols,
1531                                out_gpu,
1532                            )
1533                        });
1534                        if cpu_rows > 0 {
1535                            // Repack prefix covers the full groups of the
1536                            // CPU half (the split starts at row 0).
1537                            let rep_cpu = if repack.is_empty() {
1538                                &[][..]
1539                            } else {
1540                                &repack[..(cpu_rows / 4) * 4 * *cols]
1541                            };
1542                            qmatvec(
1543                                &bytes[..cpu_rows * *cols],
1544                                rep_cpu,
1545                                &row_scale[..cpu_rows],
1546                                x,
1547                                col_field,
1548                                *dtype,
1549                                cpu_rows,
1550                                *cols,
1551                                out_cpu,
1552                                pool,
1553                            );
1554                        }
1555                        g.join().unwrap_or(false)
1556                    });
1557                    if ok {
1558                        crate::gpu::probe_record(crate::gpu::OpClass::Matvec, true, t0.elapsed());
1559                        return;
1560                    }
1561                    // GPU failed — CPU finishes its half (rows rebased —
1562                    // group offsets don't line up, mmap layout only).
1563                    qmatvec(
1564                        &bytes[cpu_rows * *cols..(*rows) * *cols],
1565                        &[],
1566                        &row_scale[cpu_rows..],
1567                        x,
1568                        col_field,
1569                        *dtype,
1570                        *rows - cpu_rows,
1571                        *cols,
1572                        out_gpu,
1573                        pool,
1574                    );
1575                    return;
1576                }
1577                qmatvec(
1578                    self.quant_bytes(),
1579                    repack,
1580                    row_scale,
1581                    x,
1582                    col_field,
1583                    *dtype,
1584                    *rows,
1585                    *cols,
1586                    out,
1587                    pool,
1588                );
1589            }
1590        }
1591    }
1592
1593    /// Fused two-input matvec (MTP verify pair): weights streamed once.
1594    pub fn matvec2(
1595        &self,
1596        x1: &[f32],
1597        x2: &[f32],
1598        o1: &mut [f32],
1599        o2: &mut [f32],
1600        pool: Option<&Pool>,
1601    ) {
1602        match self {
1603            Self::F32 { data, .. } => matvec_rows2(pool, data, x1, x2, o1, o2),
1604            Self::Mapped {
1605                model,
1606                idx,
1607                dtype,
1608                rows,
1609                cols,
1610                row_scale,
1611                col_field,
1612                vbit_offsets,
1613                ..
1614            } => {
1615                if crate::prism::is_forward_weight(model, &model.tensors[*idx].name) {
1616                    let tx1 = crate::prism::forward(model, &x1[..*cols]);
1617                    let tx2 = crate::prism::forward(model, &x2[..*cols]);
1618                    match dtype {
1619                        TensorDtype::Q4Block => {
1620                            q4matvec2(self.quant_bytes(), &tx1, &tx2, *rows, *cols, o1, o2, pool)
1621                        }
1622                        TensorDtype::Q4Tiled => {
1623                            q4t_matvec2(self.quant_bytes(), &tx1, &tx2, *rows, *cols, o1, o2, pool)
1624                        }
1625                        TensorDtype::Q4TiledP => {
1626                            q4tp_matvec2(self.quant_bytes(), &tx1, &tx2, *rows, *cols, o1, o2, pool)
1627                        }
1628                        TensorDtype::Q2TiledP => {
1629                            if crate::prism::is_affine_target(model, &model.tensors[*idx].name) {
1630                                q2tp_affine_matvec2(
1631                                    self.quant_bytes(),
1632                                    &tx1,
1633                                    &tx2,
1634                                    *rows,
1635                                    *cols,
1636                                    o1,
1637                                    o2,
1638                                    pool,
1639                                )
1640                            } else {
1641                                q2tp_matvec2(
1642                                    self.quant_bytes(),
1643                                    &tx1,
1644                                    &tx2,
1645                                    *rows,
1646                                    *cols,
1647                                    o1,
1648                                    o2,
1649                                    pool,
1650                                )
1651                            }
1652                        }
1653                        TensorDtype::Q1 => {
1654                            q1_matvec2(self.quant_bytes(), &tx1, &tx2, *rows, *cols, o1, o2, pool)
1655                        }
1656                        TensorDtype::Q1T => {
1657                            q1t_matvec2(self.quant_bytes(), &tx1, &tx2, *rows, *cols, o1, o2, pool)
1658                        }
1659                        TensorDtype::Vbit | TensorDtype::VbitRo => vbitmatvec2(
1660                            self.quant_bytes(),
1661                            vbit_offsets,
1662                            &tx1,
1663                            &tx2,
1664                            *rows,
1665                            *cols,
1666                            o1,
1667                            o2,
1668                            pool,
1669                        ),
1670                        TensorDtype::Q8Row | TensorDtype::Q8_2f => qmatvec2(
1671                            self.quant_bytes(),
1672                            row_scale,
1673                            &tx1,
1674                            &tx2,
1675                            col_field,
1676                            *dtype,
1677                            *rows,
1678                            *cols,
1679                            o1,
1680                            o2,
1681                            pool,
1682                        ),
1683                        _ => unreachable!("unsupported mapped Prism dtype {dtype:?}"),
1684                    }
1685                    return;
1686                }
1687                if *dtype == TensorDtype::Q4Block {
1688                    q4matvec2(self.quant_bytes(), x1, x2, *rows, *cols, o1, o2, pool);
1689                    return;
1690                }
1691                if *dtype == TensorDtype::Q4Tiled {
1692                    q4t_matvec2(self.quant_bytes(), x1, x2, *rows, *cols, o1, o2, pool);
1693                    return;
1694                }
1695                if *dtype == TensorDtype::Q4TiledP {
1696                    q4tp_matvec2(self.quant_bytes(), x1, x2, *rows, *cols, o1, o2, pool);
1697                    return;
1698                }
1699                if *dtype == TensorDtype::Q2TiledP {
1700                    q2tp_matvec2(self.quant_bytes(), x1, x2, *rows, *cols, o1, o2, pool);
1701                    return;
1702                }
1703                if *dtype == TensorDtype::Q1 {
1704                    q1_matvec2(self.quant_bytes(), x1, x2, *rows, *cols, o1, o2, pool);
1705                    return;
1706                }
1707                if *dtype == TensorDtype::Q1T {
1708                    // Fused ternary pair: one row pass, the register
1709                    // unpack shared across both streams on ARM. (Q1T
1710                    // lacks a row_scale array — scales live inline in
1711                    // the tiles — so it must not fall through to the
1712                    // q8 qmatvec2 below.)
1713                    q1t_matvec2(self.quant_bytes(), x1, x2, *rows, *cols, o1, o2, pool);
1714                    return;
1715                }
1716                if matches!(dtype, TensorDtype::Vbit | TensorDtype::VbitRo) {
1717                    vbitmatvec2(
1718                        self.quant_bytes(),
1719                        vbit_offsets,
1720                        x1,
1721                        x2,
1722                        *rows,
1723                        *cols,
1724                        o1,
1725                        o2,
1726                        pool,
1727                    );
1728                    return;
1729                }
1730                qmatvec2(
1731                    self.quant_bytes(),
1732                    row_scale,
1733                    x1,
1734                    x2,
1735                    col_field,
1736                    *dtype,
1737                    *rows,
1738                    *cols,
1739                    o1,
1740                    o2,
1741                    pool,
1742                );
1743            }
1744        }
1745    }
1746}
1747
1748impl QTensor {
1749    /// Batched matvec (prefill-GEMM): xs — row-major [b, cols],
1750    /// out — row-major [b, rows]. Element-wise semantics are IDENTICAL
1751    /// to b matvec calls (same dot kernels in the same order); the win —
1752    /// the weight row streams from DRAM once per batch, not b times.
1753    /// `(model, index)` when this is a memory-mapped q4tp tensor — the
1754    /// identity a device-resident chain needs to hand `tp_matmat` the
1755    /// weight without going through this struct's own dispatch.
1756    pub fn q4tp_mapped(&self) -> Option<(&std::sync::Arc<CmfModel>, usize)> {
1757        if self.has_prism_contract() {
1758            return None;
1759        }
1760        match self {
1761            Self::Mapped {
1762                model, idx, dtype, ..
1763            } if *dtype == TensorDtype::Q4TiledP => Some((model, *idx)),
1764            _ => None,
1765        }
1766    }
1767
1768    pub fn matmat(&self, xs_all: &[f32], b: usize, out: &mut [f32], pool: Option<&Pool>) {
1769        let cols = self.cols();
1770        let rows = self.rows();
1771        debug_assert_eq!(xs_all.len(), b * cols);
1772        debug_assert_eq!(out.len(), b * rows);
1773        let _prof = crate::cpuprof::time(crate::cpuprof::Slot::Matmat);
1774        // GPTQ calibration: fold this layer's inputs into its Hessian. Only
1775        // Mapped tensors carry a directory name; the check is a relaxed
1776        // atomic load, free when not calibrating.
1777        if crate::gptq_capture::capturing() {
1778            if let Self::Mapped { model, idx, .. } = self {
1779                crate::gptq_capture::accumulate(&model.tensors[*idx].name, xs_all, b, cols);
1780            }
1781        }
1782        match self {
1783            Self::F32 { data, .. } => {
1784                let out_addr = SendMut(out.as_mut_ptr());
1785                let run = |start: usize, end: usize| {
1786                    for o in start..end {
1787                        let row = &data[o * cols..(o + 1) * cols];
1788                        for bi in 0..b {
1789                            let x = &xs_all[bi * cols..(bi + 1) * cols];
1790                            let mut acc = 0f32;
1791                            for j in 0..cols {
1792                                acc += row[j] * x[j];
1793                            }
1794                            unsafe { *out_addr.at(bi * rows + o) = acc };
1795                        }
1796                    }
1797                };
1798                dispatch_rows(pool, rows, &run);
1799            }
1800            Self::Mapped {
1801                model,
1802                idx,
1803                dtype,
1804                row_scale,
1805                col_field,
1806                vbit_offsets,
1807                ..
1808            } => {
1809                if crate::prism::is_forward_weight(model, &model.tensors[*idx].name) {
1810                    let mut transformed = Vec::with_capacity(xs_all.len());
1811                    for bi in 0..b {
1812                        transformed.extend_from_slice(&crate::prism::forward(
1813                            model,
1814                            &xs_all[bi * cols..(bi + 1) * cols],
1815                        ));
1816                    }
1817                    match dtype {
1818                        TensorDtype::Q4Block => {
1819                            q4matmat(self.quant_bytes(), &transformed, b, rows, cols, out, pool)
1820                        }
1821                        TensorDtype::Q4Tiled => {
1822                            q4t_matmat(self.quant_bytes(), &transformed, b, rows, cols, out, pool)
1823                        }
1824                        TensorDtype::Q4TiledP => {
1825                            q4tp_matmat(self.quant_bytes(), &transformed, b, rows, cols, out, pool)
1826                        }
1827                        TensorDtype::Q2TiledP => {
1828                            let affine =
1829                                crate::prism::is_affine_target(model, &model.tensors[*idx].name);
1830                            // Affine Prism Q2TP has a descriptor-aware GPU
1831                            // kernel for short/tail batches too.  Unlike the
1832                            // ordinary Q2TP path, don't force b<32 back to a
1833                            // scalar CPU matmat: prefill chunks and the final
1834                            // tail both need to stay on the tested GPU arm.
1835                            let gpu_batch_ok = if affine {
1836                                b >= 2
1837                            } else {
1838                                b >= 32 && b * rows * cols >= 128_000_000
1839                            };
1840                            if gpu_batch_ok
1841                                && cols % 32 == 0
1842                                && crate::gpu::enabled_here()
1843                                && crate::gpu::q2tp_gpu_opt_in()
1844                            {
1845                                let gpu_ok = if affine {
1846                                    crate::gpu::q2tp_affine_matmat(
1847                                        model,
1848                                        *idx,
1849                                        &transformed,
1850                                        b,
1851                                        rows,
1852                                        cols,
1853                                        out,
1854                                    )
1855                                } else {
1856                                    crate::gpu::q2tp_matmat(
1857                                        model,
1858                                        *idx,
1859                                        &transformed,
1860                                        b,
1861                                        rows,
1862                                        cols,
1863                                        out,
1864                                    )
1865                                };
1866                                if gpu_ok {
1867                                    return;
1868                                }
1869                            }
1870                            // A one-token Prism decode is the other short
1871                            // case.  Use the descriptor-aware matvec kernel
1872                            // before falling back to the exact CPU path.
1873                            if affine
1874                                && b == 1
1875                                && cols % 32 == 0
1876                                && crate::gpu::enabled_here()
1877                                && crate::gpu::q2tp_gpu_opt_in()
1878                                && crate::gpu::q2tp_affine_matvec(
1879                                    model,
1880                                    *idx,
1881                                    &transformed[..cols],
1882                                    rows,
1883                                    cols,
1884                                    &mut out[..rows],
1885                                )
1886                            {
1887                                return;
1888                            }
1889                            if affine {
1890                                q2tp_affine_matmat(
1891                                    self.quant_bytes(),
1892                                    &transformed,
1893                                    b,
1894                                    rows,
1895                                    cols,
1896                                    out,
1897                                    pool,
1898                                )
1899                            } else {
1900                                q2tp_matmat(
1901                                    self.quant_bytes(),
1902                                    &transformed,
1903                                    b,
1904                                    rows,
1905                                    cols,
1906                                    out,
1907                                    pool,
1908                                )
1909                            }
1910                        }
1911                        TensorDtype::Q1 => {
1912                            q1_matmat(self.quant_bytes(), &transformed, b, rows, cols, out, pool)
1913                        }
1914                        TensorDtype::Q1T => {
1915                            q1t_matmat(self.quant_bytes(), &transformed, b, rows, cols, out, pool)
1916                        }
1917                        TensorDtype::Vbit | TensorDtype::VbitRo => vbitmatmat(
1918                            self.quant_bytes(),
1919                            vbit_offsets,
1920                            &transformed,
1921                            b,
1922                            rows,
1923                            cols,
1924                            out,
1925                            pool,
1926                        ),
1927                        TensorDtype::Q8Row | TensorDtype::Q8_2f => {
1928                            let pre: Vec<std::borrow::Cow<'_, [f32]>> = (0..b)
1929                                .map(|bi| {
1930                                    prescale(
1931                                        &transformed[bi * cols..(bi + 1) * cols],
1932                                        col_field,
1933                                        *dtype,
1934                                    )
1935                                })
1936                                .collect();
1937                            qmatmat(self.quant_bytes(), row_scale, &pre, rows, cols, out, pool)
1938                        }
1939                        _ => unreachable!("unsupported mapped Prism dtype {dtype:?}"),
1940                    }
1941                    return;
1942                }
1943                if *dtype == TensorDtype::Q4Block {
1944                    q4matmat(self.quant_bytes(), xs_all, b, rows, cols, out, pool);
1945                    return;
1946                }
1947                if *dtype == TensorDtype::Q4TiledP {
1948                    // GPU batched q4tp GEMM (dequant + f32nt mul_mm on the
1949                    // device); the probe keeps whichever beats the CPU arm.
1950                    // Narrow (prompt-encode) and wide (DiT) batches probe
1951                    // as separate classes — the regimes have opposite
1952                    // winners and one shared verdict locked the wrong arm.
1953                    // Kill switch (gpu::mm_kill): one grossly slow GPU op
1954                    // (a fair-condition op is ≤~100 ms even at 1024px)
1955                    // means the device is contended by another process
1956                    // (e.g. a simulator) — verdicts are per-process, so
1957                    // without the bail the whole render crawls behind
1958                    // someone else's queue.
1959                    if b >= 32
1960                        && b * rows * cols >= 128_000_000
1961                        && cols % 32 == 0
1962                        && !crate::gpu::mm_killed()
1963                        && crate::gpu::enabled_here()
1964                    {
1965                        let class = if b >= 128 {
1966                            crate::gpu::OpClass::MatmatWide
1967                        } else {
1968                            crate::gpu::OpClass::Matmat
1969                        };
1970                        if let Self::Mapped { model, idx, .. } = self {
1971                            // In-process A/B (`CMF_MM_AB=1`). Three
1972                            // wall-clock A/Bs on a shared stand disagreed
1973                            // with each other by 25% on the same change,
1974                            // because the machine drifts between processes
1975                            // and interleaving whole renders does not fix
1976                            // that. Here both arms run back to back on the
1977                            // SAME data inside one call, so whatever the
1978                            // machine is doing, it does to both — and the
1979                            // disagreement between their outputs falls out
1980                            // for free. Doubles the work; a diagnostic,
1981                            // not a mode.
1982                            if crate::mm_ab::on() {
1983                                let mut g = vec![0f32; b * rows];
1984                                let t = std::time::Instant::now();
1985                                let took = crate::gpu::q4tp_matmat(
1986                                    model, *idx, xs_all, b, rows, cols, &mut g,
1987                                );
1988                                let dg = t.elapsed();
1989                                let t = std::time::Instant::now();
1990                                q4tp_matmat(self.quant_bytes(), xs_all, b, rows, cols, out, pool);
1991                                let dc = t.elapsed();
1992                                crate::mm_ab::record(b, rows, cols, took, dg, dc, &g, out);
1993                                return;
1994                            }
1995                            let t0 = std::time::Instant::now();
1996                            // A cold call takes the device arm: its sample
1997                            // is discarded either way, and the upload is
1998                            // what the next step needs.
1999                            let resident = crate::gpu::weight_is_resident(model, *idx);
2000                            match crate::gpu::probe_arm_cold_prefers_gpu(class, resident) {
2001                                crate::gpu::ProbeArm::Gpu => {
2002                                    if crate::gpu::q4tp_matmat(
2003                                        model, *idx, xs_all, b, rows, cols, out,
2004                                    ) {
2005                                        let el = t0.elapsed();
2006                                        // Work-proportional budget: ~8× the
2007                                        // fair-device estimate (+20 ms slack).
2008                                        // An absolute cap missed the worst
2009                                        // case — contended ops sit at
2010                                        // 100–240 ms each and still bury a
2011                                        // render whose fair op is 3–9 ms.
2012                                        // Cold ops (first PSO build, buffer
2013                                        // alloc) are exempt: a one-off
2014                                        // ~50 ms compile is not contention.
2015                                        let flops = 2.0 * b as f64 * rows as f64 * cols as f64;
2016                                        let budget = std::time::Duration::from_secs_f64(
2017                                            flops / 1.5e12 * 8.0 + 0.020,
2018                                        );
2019                                        crate::gpu::mm_budget_check(
2020                                            "q4tp matmat",
2021                                            el,
2022                                            budget,
2023                                            crate::gpu::probe_was_cold() || !resident,
2024                                        );
2025                                        crate::gpu::probe_record(class, true, el);
2026                                        return;
2027                                    }
2028                                }
2029                                crate::gpu::ProbeArm::CpuTimed => {
2030                                    q4tp_matmat(
2031                                        self.quant_bytes(),
2032                                        xs_all,
2033                                        b,
2034                                        rows,
2035                                        cols,
2036                                        out,
2037                                        pool,
2038                                    );
2039                                    crate::gpu::probe_record(class, false, t0.elapsed());
2040                                    return;
2041                                }
2042                                crate::gpu::ProbeArm::Cpu => {}
2043                            }
2044                        }
2045                    }
2046                    q4tp_matmat(self.quant_bytes(), xs_all, b, rows, cols, out, pool);
2047                    return;
2048                }
2049                if *dtype == TensorDtype::Q2TiledP {
2050                    // Same device arm as q4tp, behind the same probe:
2051                    // the planes differ, the dispatch does not. Without
2052                    // this a q2tp file ran its widest projections on the
2053                    // host while the 4-bit one had the card, which is a
2054                    // codec paying for its size twice.
2055                    if b >= 32
2056                        && b * rows * cols >= 128_000_000
2057                        && cols % 32 == 0
2058                        && !crate::gpu::mm_killed()
2059                        && crate::gpu::enabled_here()
2060                    {
2061                        let class = if b >= 128 {
2062                            crate::gpu::OpClass::MatmatWide
2063                        } else {
2064                            crate::gpu::OpClass::Matmat
2065                        };
2066                        if let Self::Mapped { model, idx, .. } = self {
2067                            let t0 = std::time::Instant::now();
2068                            match crate::gpu::probe_arm(class) {
2069                                crate::gpu::ProbeArm::Gpu => {
2070                                    if crate::gpu::q2tp_matmat(
2071                                        model, *idx, xs_all, b, rows, cols, out,
2072                                    ) {
2073                                        crate::gpu::probe_record(class, true, t0.elapsed());
2074                                        return;
2075                                    }
2076                                }
2077                                crate::gpu::ProbeArm::CpuTimed => {
2078                                    q2tp_matmat(
2079                                        self.quant_bytes(),
2080                                        xs_all,
2081                                        b,
2082                                        rows,
2083                                        cols,
2084                                        out,
2085                                        pool,
2086                                    );
2087                                    crate::gpu::probe_record(class, false, t0.elapsed());
2088                                    return;
2089                                }
2090                                crate::gpu::ProbeArm::Cpu => {}
2091                            }
2092                        }
2093                    }
2094                    // Without a host arm a q2tp tensor falls through to
2095                    // the q8 fallback, which reads it at one BYTE per
2096                    // weight — a 2x overrun that killed pool workers
2097                    // mid-prefill while the dispatcher waited forever.
2098                    q2tp_matmat(self.quant_bytes(), xs_all, b, rows, cols, out, pool);
2099                    return;
2100                }
2101                if *dtype == TensorDtype::Q4Tiled {
2102                    // GPU batched q4t GEMM (dequant + f32nt mul_mm on the
2103                    // device); the probe keeps whichever beats the CPU arm.
2104                    // Narrow (prompt-encode) and wide (DiT) batches probe
2105                    // as separate classes — the regimes have opposite
2106                    // winners and one shared verdict locked the wrong arm.
2107                    // Kill switch (gpu::mm_kill): one grossly slow GPU op
2108                    // (a fair-condition op is ≤~100 ms even at 1024px)
2109                    // means the device is contended by another process
2110                    // (e.g. a simulator) — verdicts are per-process, so
2111                    // without the bail the whole render crawls behind
2112                    // someone else's queue.
2113                    if b >= 32
2114                        && b * rows * cols >= 128_000_000
2115                        && cols % 32 == 0
2116                        && !crate::gpu::mm_killed()
2117                        && crate::gpu::enabled_here()
2118                    {
2119                        let class = if b >= 128 {
2120                            crate::gpu::OpClass::MatmatWide
2121                        } else {
2122                            crate::gpu::OpClass::Matmat
2123                        };
2124                        if let Self::Mapped { model, idx, .. } = self {
2125                            let t0 = std::time::Instant::now();
2126                            match crate::gpu::probe_arm(class) {
2127                                crate::gpu::ProbeArm::Gpu => {
2128                                    if crate::gpu::q4t_matmat(
2129                                        model, *idx, xs_all, b, rows, cols, out,
2130                                    ) {
2131                                        let el = t0.elapsed();
2132                                        // Work-proportional budget: ~8× the
2133                                        // fair-device estimate (+20 ms slack).
2134                                        // An absolute cap missed the worst
2135                                        // case — contended ops sit at
2136                                        // 100–240 ms each and still bury a
2137                                        // render whose fair op is 3–9 ms.
2138                                        // Cold ops (first PSO build, buffer
2139                                        // alloc) are exempt: a one-off
2140                                        // ~50 ms compile is not contention.
2141                                        let flops = 2.0 * b as f64 * rows as f64 * cols as f64;
2142                                        let budget = std::time::Duration::from_secs_f64(
2143                                            flops / 1.5e12 * 8.0 + 0.020,
2144                                        );
2145                                        crate::gpu::mm_budget_check(
2146                                            "q4t matmat",
2147                                            el,
2148                                            budget,
2149                                            crate::gpu::probe_was_cold(),
2150                                        );
2151                                        crate::gpu::probe_record(class, true, el);
2152                                        return;
2153                                    }
2154                                }
2155                                crate::gpu::ProbeArm::CpuTimed => {
2156                                    q4t_matmat(
2157                                        self.quant_bytes(),
2158                                        xs_all,
2159                                        b,
2160                                        rows,
2161                                        cols,
2162                                        out,
2163                                        pool,
2164                                    );
2165                                    crate::gpu::probe_record(class, false, t0.elapsed());
2166                                    return;
2167                                }
2168                                crate::gpu::ProbeArm::Cpu => {}
2169                            }
2170                        }
2171                    }
2172                    q4t_matmat(self.quant_bytes(), xs_all, b, rows, cols, out, pool);
2173                    return;
2174                }
2175                if *dtype == TensorDtype::Q1 {
2176                    // GPU batched q1 GEMM for wide prefill (q1_mul_mm on the
2177                    // device); the probe keeps whichever beats the CPU matmat.
2178                    if b >= 32
2179                        && b * rows * cols >= 128_000_000
2180                        && cols % 64 == 0
2181                        && crate::gpu::enabled_here()
2182                    {
2183                        if let Self::Mapped { model, idx, .. } = self {
2184                            let t0 = std::time::Instant::now();
2185                            match crate::gpu::probe_arm(crate::gpu::OpClass::Matmat) {
2186                                crate::gpu::ProbeArm::Gpu => {
2187                                    if crate::gpu::q1_matmat(
2188                                        model, *idx, xs_all, b, rows, cols, out,
2189                                    ) {
2190                                        crate::gpu::probe_record(
2191                                            crate::gpu::OpClass::Matmat,
2192                                            true,
2193                                            t0.elapsed(),
2194                                        );
2195                                        return;
2196                                    }
2197                                }
2198                                crate::gpu::ProbeArm::CpuTimed => {
2199                                    q1_matmat(self.quant_bytes(), xs_all, b, rows, cols, out, pool);
2200                                    crate::gpu::probe_record(
2201                                        crate::gpu::OpClass::Matmat,
2202                                        false,
2203                                        t0.elapsed(),
2204                                    );
2205                                    return;
2206                                }
2207                                crate::gpu::ProbeArm::Cpu => {}
2208                            }
2209                        }
2210                    }
2211                    q1_matmat(self.quant_bytes(), xs_all, b, rows, cols, out, pool);
2212                    return;
2213                }
2214                if *dtype == TensorDtype::Q1T {
2215                    // GPU batched GEMM for wide prefill (base + overlay on the
2216                    // device); probe keeps the winner vs the CPU matmat.
2217                    if b >= 32 && b * rows * cols >= 128_000_000 && crate::gpu::enabled_here() {
2218                        if let Self::Mapped { model, idx, .. } = self {
2219                            let t0 = std::time::Instant::now();
2220                            match crate::gpu::probe_arm(crate::gpu::OpClass::Matmat) {
2221                                crate::gpu::ProbeArm::Gpu => {
2222                                    if crate::gpu::q1t_matmat(
2223                                        model, *idx, xs_all, b, rows, cols, out,
2224                                    ) {
2225                                        crate::gpu::probe_record(
2226                                            crate::gpu::OpClass::Matmat,
2227                                            true,
2228                                            t0.elapsed(),
2229                                        );
2230                                        return;
2231                                    }
2232                                }
2233                                crate::gpu::ProbeArm::CpuTimed => {
2234                                    q1t_matmat(
2235                                        self.quant_bytes(),
2236                                        xs_all,
2237                                        b,
2238                                        rows,
2239                                        cols,
2240                                        out,
2241                                        pool,
2242                                    );
2243                                    crate::gpu::probe_record(
2244                                        crate::gpu::OpClass::Matmat,
2245                                        false,
2246                                        t0.elapsed(),
2247                                    );
2248                                    return;
2249                                }
2250                                crate::gpu::ProbeArm::Cpu => {}
2251                            }
2252                        }
2253                    }
2254                    q1t_matmat(self.quant_bytes(), xs_all, b, rows, cols, out, pool);
2255                    return;
2256                }
2257                if matches!(dtype, TensorDtype::Vbit | TensorDtype::VbitRo) {
2258                    vbitmatmat(
2259                        self.quant_bytes(),
2260                        vbit_offsets,
2261                        xs_all,
2262                        b,
2263                        rows,
2264                        cols,
2265                        out,
2266                        pool,
2267                    );
2268                    return;
2269                }
2270                let pre: Vec<std::borrow::Cow<'_, [f32]>> = (0..b)
2271                    .map(|bi| prescale(&xs_all[bi * cols..(bi + 1) * cols], col_field, *dtype))
2272                    .collect();
2273                // MiMo verification is a 2–4 row decode panel, not a wide
2274                // prompt GEMM. Keep q8 projections on the same device as
2275                // decode; the generic b>=8 gate otherwise silently moves
2276                // every projection back to CPU. The short wgpu matmat uses
2277                // the same 64-lane reduction as its single-token matvec.
2278                if row_exact()
2279                    && (1..=4).contains(&b)
2280                    && matches!(dtype, TensorDtype::Q8Row | TensorDtype::Q8_2f)
2281                    && crate::gpu::enabled_here()
2282                    && crate::gpu::wgpu_active()
2283                {
2284                    let flat: Vec<f32> = pre.iter().flat_map(|v| v.iter().copied()).collect();
2285                    if crate::gpu::q8_matmat(model, *idx, row_scale, &flat, b, rows, cols, out) {
2286                        return;
2287                    }
2288                }
2289                // D5: large prefill-batch GEMMs — on the GPU (threshold by
2290                // work volume: submission carries b×rows×cols MACs).
2291                // Runtime probe: the naive GEMM shader + sync readback
2292                // lose to the CPU GEMM on slow driver stacks — alternate
2293                // both arms and keep the winner.
2294                if b >= 8 && b * rows * cols >= 128_000_000 && crate::gpu::enabled_here() {
2295                    if let Self::Mapped { model, idx, .. } = self {
2296                        let t0 = std::time::Instant::now();
2297                        match crate::gpu::probe_arm(crate::gpu::OpClass::Matmat) {
2298                            crate::gpu::ProbeArm::Gpu
2299                                if crate::gpu::probe_deciding(crate::gpu::OpClass::Matmat)
2300                                    && !crate::gpu::q8_resident_or_upload(model, *idx) =>
2301                            {
2302                                // Cold weights during probing: the upload
2303                                // has started, the count runs on the CPU —
2304                                // the GPU arm samples on the next touch.
2305                                let q = self.quant_bytes();
2306                                qmatmat(q, row_scale, &pre, rows, cols, out, pool);
2307                                return;
2308                            }
2309                            crate::gpu::ProbeArm::Gpu => {
2310                                let flat: Vec<f32> =
2311                                    pre.iter().flat_map(|v| v.iter().copied()).collect();
2312                                if crate::gpu::q8_matmat(
2313                                    model, *idx, row_scale, &flat, b, rows, cols, out,
2314                                ) {
2315                                    crate::gpu::probe_record(
2316                                        crate::gpu::OpClass::Matmat,
2317                                        true,
2318                                        t0.elapsed(),
2319                                    );
2320                                    return;
2321                                }
2322                            }
2323                            crate::gpu::ProbeArm::CpuTimed => {
2324                                let q = self.quant_bytes();
2325                                qmatmat(q, row_scale, &pre, rows, cols, out, pool);
2326                                crate::gpu::probe_record(
2327                                    crate::gpu::OpClass::Matmat,
2328                                    false,
2329                                    t0.elapsed(),
2330                                );
2331                                return;
2332                            }
2333                            crate::gpu::ProbeArm::Cpu => {}
2334                        }
2335                    }
2336                }
2337                let q = self.quant_bytes();
2338                qmatmat(q, row_scale, &pre, rows, cols, out, pool);
2339            }
2340        }
2341    }
2342}
2343
2344impl QTensor {
2345    /// The device GEMM this tensor would take, run once on the caller's
2346    /// data — the startup parity probe's arm, and the one place that knows
2347    /// which entry point each codec has.
2348    ///
2349    /// It exists because the probe used to look for a `q4tp` weight by
2350    /// name AND dtype, and a container packed any other way was declared
2351    /// "host path" for the whole render even though its codec had a device
2352    /// GEMM of its own. A gate that only recognizes one codec is a gate
2353    /// that silently downgrades every other one.
2354    pub fn device_matmat(&self, xs: &[f32], b: usize, out: &mut [f32]) -> bool {
2355        let (rows, cols) = (self.rows(), self.cols());
2356        let Self::Mapped {
2357            model,
2358            idx,
2359            dtype,
2360            row_scale,
2361            col_field,
2362            ..
2363        } = self
2364        else {
2365            return false;
2366        };
2367        if crate::prism::has_contract(model) {
2368            return false;
2369        }
2370        match *dtype {
2371            TensorDtype::Q4TiledP => crate::gpu::q4tp_matmat(model, *idx, xs, b, rows, cols, out),
2372            // The two-field codec folds its column field into the
2373            // activation, which leaves a plain per-row int8 GEMM — the
2374            // same kernel `q8_row` uses, on both backends.
2375            TensorDtype::Q8Row | TensorDtype::Q8_2f => {
2376                // The field belongs to the weight; only a backend that cannot
2377                // apply it there makes a scaled copy of the activation.
2378                if *dtype == TensorDtype::Q8_2f
2379                    && std::env::var("CMF_Q8_2F_DEV").as_deref() != Ok("0")
2380                    && crate::gpu::q8_matmat_2f(
2381                        model, *idx, row_scale, col_field, xs, b, rows, cols, out,
2382                    )
2383                {
2384                    return true;
2385                }
2386                let flat: Vec<f32> = (0..b)
2387                    .flat_map(|bi| {
2388                        prescale(&xs[bi * cols..(bi + 1) * cols], col_field, *dtype).into_owned()
2389                    })
2390                    .collect();
2391                crate::gpu::q8_matmat(model, *idx, row_scale, &flat, b, rows, cols, out)
2392            }
2393            _ => false,
2394        }
2395    }
2396
2397    /// Multi-matrix job (roadmap §3 P0): N tensors sharing one input
2398    /// run under a SINGLE pool dispatch — QKV or gate+up cost one
2399    /// barrier instead of N. Per-row math is the exact same kernel as
2400    /// `matvec` (bit-identical outputs); only the dispatch is fused.
2401    /// Falls back to N sequential matvecs when the set is not a uniform
2402    /// q8-family/F32 group or there is no pool.
2403    pub fn matvec_many<const N: usize>(
2404        ts: [&QTensor; N],
2405        x: &[f32],
2406        mut outs: [&mut [f32]; N],
2407        pool: Option<&Pool>,
2408    ) {
2409        let total_rows: usize = ts.iter().map(|t| t.rows()).sum();
2410        if ts.iter().any(|t| t.has_prism_contract()) {
2411            // The fused range kernels have no transform descriptor.  Let
2412            // each tensor's ordinary matvec dispatch perform the explicit
2413            // signed FWHT (and retain CPU fallback for mixed q2tp/q4tp).
2414            for (t, o) in ts.iter().zip(outs.iter_mut()) {
2415                t.matvec(x, o, pool);
2416            }
2417            return;
2418        }
2419        let uniform_q8 = ts.iter().all(|t| {
2420            matches!(
2421                t,
2422                Self::Mapped {
2423                    dtype: TensorDtype::Q8Row | TensorDtype::Q8_2f,
2424                    ..
2425                }
2426            )
2427        });
2428        let uniform_f32 = ts.iter().all(|t| matches!(t, Self::F32 { .. }));
2429        let uniform_q4 = ts.iter().all(|t| {
2430            matches!(
2431                t,
2432                Self::Mapped {
2433                    dtype: TensorDtype::Q4Block,
2434                    ..
2435                }
2436            )
2437        });
2438        let uniform_vbit = ts.iter().all(|t| {
2439            matches!(
2440                t,
2441                Self::Mapped {
2442                    dtype: TensorDtype::Vbit | TensorDtype::VbitRo,
2443                    ..
2444                }
2445            )
2446        });
2447        let uniform_q1 = ts.iter().all(|t| {
2448            matches!(
2449                t,
2450                Self::Mapped {
2451                    dtype: TensorDtype::Q1,
2452                    ..
2453                }
2454            )
2455        });
2456        let uniform_q1t = ts.iter().all(|t| {
2457            matches!(
2458                t,
2459                Self::Mapped {
2460                    dtype: TensorDtype::Q1T,
2461                    ..
2462                }
2463            )
2464        });
2465        // q4tp is the skeleton dtype of the big MoE files, and without an arm
2466        // here every projection that shares an input paid its own pool
2467        // barrier: DeepSeek-V4's attention step alone hands this function
2468        // wq_a, wkv and both compressors' pairs off the same hidden state.
2469        let uniform_q4tp = ts.iter().all(|t| {
2470            matches!(
2471                t,
2472                Self::Mapped {
2473                    dtype: TensorDtype::Q4TiledP,
2474                    ..
2475                }
2476            )
2477        }) && ts
2478            .iter()
2479            .all(|t| t.cols() == ts[0].cols() && t.cols() % GROUP_SIZE == 0);
2480        let Some(pool) = pool else {
2481            for (t, o) in ts.iter().zip(outs.iter_mut()) {
2482                t.matvec(x, o, None);
2483            }
2484            return;
2485        };
2486        if total_rows < 256
2487            || !(uniform_q8
2488                || uniform_f32
2489                || uniform_q4
2490                || uniform_vbit
2491                || uniform_q1
2492                || uniform_q1t
2493                || uniform_q4tp)
2494        {
2495            for (t, o) in ts.iter().zip(outs.iter_mut()) {
2496                t.matvec(x, o, Some(pool));
2497            }
2498            return;
2499        }
2500
2501        if uniform_q4tp {
2502            // Every tensor's rows laid end to end in one virtual row space,
2503            // so the whole set is ONE dispatch. The per-row body is the
2504            // `q4tp_matvec` arm verbatim — same activation split, same
2505            // accumulation order — so the outputs are bit-identical to the
2506            // sequential calls this replaces.
2507            let cols = ts[0].cols();
2508            let gpr = cols / GROUP_SIZE;
2509            let views: Vec<Q4tpView> = ts
2510                .iter()
2511                .map(|t| Q4tpView::new(t.quant_bytes(), t.rows(), cols))
2512                .collect();
2513            let rows_of: Vec<usize> = ts.iter().map(|t| t.rows()).collect();
2514            let outs_addr: [SendMut; N] = std::array::from_fn(|i| SendMut(outs[i].as_mut_ptr()));
2515            // flat index -> (which tensor, which of its rows)
2516            let locate = |flat: usize| -> (usize, usize) {
2517                let mut acc = 0;
2518                for (i, &r) in rows_of.iter().enumerate() {
2519                    if flat < acc + r {
2520                        return (i, flat - acc);
2521                    }
2522                    acc += r;
2523                }
2524                (rows_of.len() - 1, 0)
2525            };
2526            let (views, outs_addr) = (&views, &outs_addr);
2527            if a8w8_enabled() {
2528                let act = split_act(x);
2529                let act = &act;
2530                let run = |start: usize, end: usize| {
2531                    let mut sc = vec![0f32; gpr];
2532                    for flat in start..end {
2533                        let (t, r) = locate(flat);
2534                        let v = &views[t];
2535                        v.scales_into(r, gpr, &mut sc);
2536                        let mut acc = dot_q4tp_row_i8(v.nib, r, gpr, &act.xq, &sc) * act.sx;
2537                        for &(j, xv) in &act.outliers {
2538                            let (w, s) = q4tp_outlier(v.nib, r, gpr, j, &sc);
2539                            acc += w * s * xv;
2540                        }
2541                        // SAFETY: one worker owns each (tensor, row) pair.
2542                        unsafe { *outs_addr[t].at(r) = acc };
2543                    }
2544                };
2545                pool.run_rows(total_rows, &run);
2546            } else {
2547                let run = |start: usize, end: usize| {
2548                    let mut sc = vec![0f32; gpr];
2549                    for flat in start..end {
2550                        let (t, r) = locate(flat);
2551                        let v = &views[t];
2552                        v.scales_into(r, gpr, &mut sc);
2553                        // SAFETY: one worker owns each (tensor, row) pair.
2554                        unsafe { *outs_addr[t].at(r) = q4tp_row_exact(v.nib, r, gpr, x, &sc) };
2555                    }
2556                };
2557                pool.run_rows(total_rows, &run);
2558            }
2559            return;
2560        }
2561
2562        if uniform_q1 {
2563            // One shared activation split + group sums (q1 has no col
2564            // field; the same input feeds every tensor).
2565            let outs_addr: [SendMut; N] = std::array::from_fn(|i| SendMut(outs[i].as_mut_ptr()));
2566            if a8w8_enabled() {
2567                let act = split_act(x);
2568                let gsum = q1_group_sums(&act.xq, ts[0].cols() / GROUP_SIZE);
2569                let (act, gsum) = (&act, &gsum);
2570                let closures: [_; N] = std::array::from_fn(|i| {
2571                    let (bytes, gpr, out) =
2572                        (ts[i].quant_bytes(), ts[i].cols() / GROUP_SIZE, outs_addr[i]);
2573                    move |s: usize, e: usize| q1_range_a8w8(bytes, gpr, act, gsum, out, s, e)
2574                });
2575                let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
2576                    std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
2577                pool.run_many(&parts);
2578            } else {
2579                let closures: [_; N] = std::array::from_fn(|i| {
2580                    let (bytes, gpr, out) =
2581                        (ts[i].quant_bytes(), ts[i].cols() / GROUP_SIZE, outs_addr[i]);
2582                    move |s: usize, e: usize| q1_range_f32(bytes, gpr, x, out, s, e)
2583                });
2584                let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
2585                    std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
2586                pool.run_many(&parts);
2587            }
2588            return;
2589        }
2590
2591        if uniform_q1t {
2592            // Q1T batched: one shared activation split + overlay decode,
2593            // all tensors' rows in ONE pool dispatch (saves N−1 dispatches
2594            // and N−1 redundant split_act calls per layer).
2595            let outs_addr: [SendMut; N] = std::array::from_fn(|i| SendMut(outs[i].as_mut_ptr()));
2596            const TILE: usize = cortiq_core::quant::Q1T_TILE;
2597            if a8w8_enabled() {
2598                let act = split_act(x);
2599                let act = &act;
2600                let x_ref = x;
2601                let closures: [_; N] = std::array::from_fn(|i| {
2602                    let bytes = ts[i].quant_bytes();
2603                    let (rows, cols) = (ts[i].rows(), ts[i].cols());
2604                    let gpr = cols / GROUP_SIZE;
2605                    let (rp_off, ent_off, has_ov) = q1t_overlay(bytes, rows * gpr * TILE, rows);
2606                    let out = outs_addr[i];
2607                    move |s: usize, e: usize| {
2608                        q1t_range_a8w8(bytes, gpr, rp_off, ent_off, has_ov, act, x_ref, out, s, e)
2609                    }
2610                });
2611                let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
2612                    std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
2613                pool.run_many(&parts);
2614            } else {
2615                let x_ref = x;
2616                let closures: [_; N] = std::array::from_fn(|i| {
2617                    let bytes = ts[i].quant_bytes();
2618                    let (rows, cols) = (ts[i].rows(), ts[i].cols());
2619                    let gpr = cols / GROUP_SIZE;
2620                    let (rp_off, ent_off, has_ov) = q1t_overlay(bytes, rows * gpr * TILE, rows);
2621                    let out = outs_addr[i];
2622                    move |s: usize, e: usize| {
2623                        q1t_range_f32_batch(bytes, gpr, rp_off, ent_off, has_ov, x_ref, out, s, e)
2624                    }
2625                });
2626                let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
2627                    std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
2628                pool.run_many(&parts);
2629            }
2630            return;
2631        }
2632
2633        if uniform_q4 || uniform_vbit {
2634            let outs_addr: [SendMut; N] = std::array::from_fn(|i| SendMut(outs[i].as_mut_ptr()));
2635            // q4/vbit share one activation split — no per-tensor col field.
2636            if a8w8_enabled() {
2637                let act = split_act(x);
2638                let act = &act;
2639                if uniform_q4 {
2640                    let closures: [_; N] = std::array::from_fn(|i| {
2641                        let (packed, scales) =
2642                            q4_split(ts[i].quant_bytes(), ts[i].rows(), ts[i].cols());
2643                        let (gpr, cols, out) =
2644                            (ts[i].cols() / GROUP_SIZE, ts[i].cols(), outs_addr[i]);
2645                        move |s: usize, e: usize| {
2646                            q4_range_a8w8(packed, scales, gpr, cols, act, out, s, e)
2647                        }
2648                    });
2649                    let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
2650                        std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
2651                    pool.run_many(&parts);
2652                } else {
2653                    let closures: [_; N] = std::array::from_fn(|i| {
2654                        let Self::Mapped { vbit_offsets, .. } = ts[i] else {
2655                            unreachable!()
2656                        };
2657                        let (bytes, rows, cols, out) = (
2658                            ts[i].quant_bytes(),
2659                            ts[i].rows(),
2660                            ts[i].cols(),
2661                            outs_addr[i],
2662                        );
2663                        move |s: usize, e: usize| {
2664                            vbit_range_a8w8(bytes, vbit_offsets, x, act, rows, cols, out, s, e)
2665                        }
2666                    });
2667                    let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
2668                        std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
2669                    pool.run_many(&parts);
2670                }
2671                return;
2672            }
2673            if uniform_q4 {
2674                let closures: [_; N] = std::array::from_fn(|i| {
2675                    let (packed, scales) =
2676                        q4_split(ts[i].quant_bytes(), ts[i].rows(), ts[i].cols());
2677                    let (gpr, out) = (ts[i].cols() / GROUP_SIZE, outs_addr[i]);
2678                    move |s: usize, e: usize| q4_range_f32(packed, scales, gpr, x, out, s, e)
2679                });
2680                let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
2681                    std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
2682                pool.run_many(&parts);
2683            } else {
2684                let closures: [_; N] = std::array::from_fn(|i| {
2685                    let Self::Mapped { vbit_offsets, .. } = ts[i] else {
2686                        unreachable!()
2687                    };
2688                    let (bytes, rows, cols, out) = (
2689                        ts[i].quant_bytes(),
2690                        ts[i].rows(),
2691                        ts[i].cols(),
2692                        outs_addr[i],
2693                    );
2694                    move |s: usize, e: usize| {
2695                        vbit_range_f32(bytes, vbit_offsets, x, rows, cols, out, s, e)
2696                    }
2697                });
2698                let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
2699                    std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
2700                pool.run_many(&parts);
2701            }
2702            return;
2703        }
2704
2705        if uniform_f32 {
2706            let outs_addr: [SendMut; N] = std::array::from_fn(|i| SendMut(outs[i].as_mut_ptr()));
2707            let closures: [_; N] = std::array::from_fn(|i| {
2708                let Self::F32 { data, cols, .. } = ts[i] else {
2709                    unreachable!()
2710                };
2711                let out = outs_addr[i];
2712                move |start: usize, end: usize| {
2713                    for o in start..end {
2714                        let row = &data[o * cols..(o + 1) * cols];
2715                        let mut sum = 0.0f32;
2716                        for j in 0..*cols {
2717                            sum += row[j] * x[j];
2718                        }
2719                        // SAFETY: disjoint (tensor, row) cells per worker.
2720                        unsafe { *out.at(o) = sum };
2721                    }
2722                }
2723            });
2724            let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
2725                std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
2726            pool.run_many(&parts);
2727            return;
2728        }
2729
2730        // Uniform q8-family: per-tensor prescale (q8_2f col fields
2731        // differ per tensor) + the shared range kernels.
2732        struct Ctx<'a> {
2733            bytes: &'a [u8],
2734            #[cfg_attr(not(target_arch = "aarch64"), allow(dead_code))]
2735            rep: &'a [u8],
2736            row_scale: &'a [f32],
2737            cols: usize,
2738            xs: std::borrow::Cow<'a, [f32]>,
2739        }
2740        let ctxs: [Ctx<'_>; N] = std::array::from_fn(|i| {
2741            let Self::Mapped {
2742                dtype,
2743                cols,
2744                row_scale,
2745                col_field,
2746                repack,
2747                ..
2748            } = ts[i]
2749            else {
2750                unreachable!()
2751            };
2752            Ctx {
2753                bytes: ts[i].quant_bytes(),
2754                rep: repack,
2755                row_scale,
2756                cols: *cols,
2757                xs: prescale(x, col_field, *dtype),
2758            }
2759        });
2760        let outs_addr: [SendMut; N] = std::array::from_fn(|i| SendMut(outs[i].as_mut_ptr()));
2761        #[cfg(target_arch = "aarch64")]
2762        if sdot_enabled() {
2763            let acts: [SplitAct; N] = std::array::from_fn(|i| split_act(&ctxs[i].xs));
2764            let closures: [_; N] = std::array::from_fn(|i| {
2765                let (c, act, out) = (&ctxs[i], &acts[i], outs_addr[i]);
2766                move |start: usize, end: usize| {
2767                    q8_range_sdot(c.bytes, c.rep, c.row_scale, act, c.cols, out, start, end)
2768                }
2769            });
2770            let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
2771                std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
2772            pool.run_many(&parts);
2773            return;
2774        }
2775        #[cfg(target_arch = "x86_64")]
2776        if avx2_a8w8_enabled() {
2777            let acts: [SplitAct; N] = std::array::from_fn(|i| split_act(&ctxs[i].xs));
2778            let closures: [_; N] = std::array::from_fn(|i| {
2779                let (c, act, out) = (&ctxs[i], &acts[i], outs_addr[i]);
2780                move |start: usize, end: usize| {
2781                    q8_range_avx2(c.bytes, c.row_scale, act, c.cols, out, start, end)
2782                }
2783            });
2784            let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
2785                std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
2786            pool.run_many(&parts);
2787            return;
2788        }
2789        let closures: [_; N] = std::array::from_fn(|i| {
2790            let (c, out) = (&ctxs[i], outs_addr[i]);
2791            move |start: usize, end: usize| {
2792                q8_range_f32(c.bytes, c.row_scale, &c.xs, c.cols, out, start, end)
2793            }
2794        });
2795        let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
2796            std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
2797        pool.run_many(&parts);
2798    }
2799}
2800
2801impl QTensor {
2802    /// Pair-input multi-matrix job: N tensors × 2 shared inputs under a
2803    /// single pool dispatch — the MTP/pair decode path publishes one job
2804    /// for Q/K/V (and one for gate+up) instead of one per tensor.
2805    /// Per-row math is exactly `matvec2`'s kernels; bit-identical.
2806    #[allow(clippy::needless_range_loop)]
2807    pub fn matvec2_many<const N: usize>(
2808        ts: [&QTensor; N],
2809        x1: &[f32],
2810        x2: &[f32],
2811        mut o1s: [&mut [f32]; N],
2812        mut o2s: [&mut [f32]; N],
2813        pool: Option<&Pool>,
2814    ) {
2815        let total_rows: usize = ts.iter().map(|t| t.rows()).sum();
2816        if ts.iter().any(|t| t.has_prism_contract()) {
2817            for i in 0..N {
2818                ts[i].matvec2(x1, x2, o1s[i], o2s[i], pool);
2819            }
2820            return;
2821        }
2822        let uniform_q8 = ts.iter().all(|t| {
2823            matches!(
2824                t,
2825                Self::Mapped {
2826                    dtype: TensorDtype::Q8Row | TensorDtype::Q8_2f,
2827                    ..
2828                }
2829            )
2830        });
2831        let uniform_f32 = ts.iter().all(|t| matches!(t, Self::F32 { .. }));
2832        let uniform_q4 = ts.iter().all(|t| {
2833            matches!(
2834                t,
2835                Self::Mapped {
2836                    dtype: TensorDtype::Q4Block,
2837                    ..
2838                }
2839            )
2840        });
2841        let uniform_vbit = ts.iter().all(|t| {
2842            matches!(
2843                t,
2844                Self::Mapped {
2845                    dtype: TensorDtype::Vbit | TensorDtype::VbitRo,
2846                    ..
2847                }
2848            )
2849        });
2850        let fusable = pool.is_some()
2851            && total_rows >= 256
2852            && (uniform_q8 || uniform_f32 || uniform_q4 || uniform_vbit);
2853        if !fusable {
2854            for i in 0..N {
2855                ts[i].matvec2(x1, x2, o1s[i], o2s[i], pool);
2856            }
2857            return;
2858        }
2859        let pool = pool.unwrap();
2860
2861        if uniform_q4 || uniform_vbit {
2862            let p1: [SendMut; N] = std::array::from_fn(|i| SendMut(o1s[i].as_mut_ptr()));
2863            let p2: [SendMut; N] = std::array::from_fn(|i| SendMut(o2s[i].as_mut_ptr()));
2864            // q4/vbit share activation splits — no per-tensor col field.
2865            if a8w8_enabled() {
2866                let a1 = split_act(x1);
2867                let a2 = split_act(x2);
2868                let (a1, a2) = (&a1, &a2);
2869                if uniform_q4 {
2870                    let closures: [_; N] = std::array::from_fn(|i| {
2871                        let (packed, scales) =
2872                            q4_split(ts[i].quant_bytes(), ts[i].rows(), ts[i].cols());
2873                        let (gpr, cols, o1, o2) =
2874                            (ts[i].cols() / GROUP_SIZE, ts[i].cols(), p1[i], p2[i]);
2875                        move |s: usize, e: usize| {
2876                            q4_range2_a8w8(packed, scales, gpr, cols, a1, a2, o1, o2, s, e)
2877                        }
2878                    });
2879                    let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
2880                        std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
2881                    pool.run_many(&parts);
2882                } else {
2883                    let closures: [_; N] = std::array::from_fn(|i| {
2884                        let Self::Mapped { vbit_offsets, .. } = ts[i] else {
2885                            unreachable!()
2886                        };
2887                        let (bytes, rows, cols, o1, o2) = (
2888                            ts[i].quant_bytes(),
2889                            ts[i].rows(),
2890                            ts[i].cols(),
2891                            p1[i],
2892                            p2[i],
2893                        );
2894                        move |s: usize, e: usize| {
2895                            vbit_range2_a8w8(
2896                                bytes,
2897                                vbit_offsets,
2898                                x1,
2899                                x2,
2900                                a1,
2901                                a2,
2902                                rows,
2903                                cols,
2904                                o1,
2905                                o2,
2906                                s,
2907                                e,
2908                            )
2909                        }
2910                    });
2911                    let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
2912                        std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
2913                    pool.run_many(&parts);
2914                }
2915                return;
2916            }
2917            if uniform_q4 {
2918                let closures: [_; N] = std::array::from_fn(|i| {
2919                    let (packed, scales) =
2920                        q4_split(ts[i].quant_bytes(), ts[i].rows(), ts[i].cols());
2921                    let (gpr, o1, o2) = (ts[i].cols() / GROUP_SIZE, p1[i], p2[i]);
2922                    move |s: usize, e: usize| {
2923                        q4_range2_f32(packed, scales, gpr, x1, x2, o1, o2, s, e)
2924                    }
2925                });
2926                let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
2927                    std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
2928                pool.run_many(&parts);
2929            } else {
2930                let closures: [_; N] = std::array::from_fn(|i| {
2931                    let Self::Mapped { vbit_offsets, .. } = ts[i] else {
2932                        unreachable!()
2933                    };
2934                    let (bytes, rows, cols, o1, o2) = (
2935                        ts[i].quant_bytes(),
2936                        ts[i].rows(),
2937                        ts[i].cols(),
2938                        p1[i],
2939                        p2[i],
2940                    );
2941                    move |s: usize, e: usize| {
2942                        vbit_range2_f32(bytes, vbit_offsets, x1, x2, rows, cols, o1, o2, s, e)
2943                    }
2944                });
2945                let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
2946                    std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
2947                pool.run_many(&parts);
2948            }
2949            return;
2950        }
2951
2952        if uniform_f32 {
2953            let p1: [SendMut; N] = std::array::from_fn(|i| SendMut(o1s[i].as_mut_ptr()));
2954            let p2: [SendMut; N] = std::array::from_fn(|i| SendMut(o2s[i].as_mut_ptr()));
2955            let closures: [_; N] = std::array::from_fn(|i| {
2956                let Self::F32 { data, cols, .. } = ts[i] else {
2957                    unreachable!()
2958                };
2959                let (o1, o2) = (p1[i], p2[i]);
2960                move |start: usize, end: usize| {
2961                    for o in start..end {
2962                        let row = &data[o * cols..(o + 1) * cols];
2963                        let (mut s1, mut s2) = (0.0f32, 0.0f32);
2964                        for j in 0..*cols {
2965                            s1 += row[j] * x1[j];
2966                            s2 += row[j] * x2[j];
2967                        }
2968                        // SAFETY: disjoint (tensor, row) cells per worker.
2969                        unsafe {
2970                            *o1.at(o) = s1;
2971                            *o2.at(o) = s2;
2972                        }
2973                    }
2974                }
2975            });
2976            let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
2977                std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
2978            pool.run_many(&parts);
2979            return;
2980        }
2981
2982        struct Ctx<'a> {
2983            bytes: &'a [u8],
2984            row_scale: &'a [f32],
2985            cols: usize,
2986            xs1: std::borrow::Cow<'a, [f32]>,
2987            xs2: std::borrow::Cow<'a, [f32]>,
2988        }
2989        let ctxs: [Ctx<'_>; N] = std::array::from_fn(|i| {
2990            let Self::Mapped {
2991                dtype,
2992                cols,
2993                row_scale,
2994                col_field,
2995                ..
2996            } = ts[i]
2997            else {
2998                unreachable!()
2999            };
3000            Ctx {
3001                bytes: ts[i].quant_bytes(),
3002                row_scale,
3003                cols: *cols,
3004                xs1: prescale(x1, col_field, *dtype),
3005                xs2: prescale(x2, col_field, *dtype),
3006            }
3007        });
3008        let p1: [SendMut; N] = std::array::from_fn(|i| SendMut(o1s[i].as_mut_ptr()));
3009        let p2: [SendMut; N] = std::array::from_fn(|i| SendMut(o2s[i].as_mut_ptr()));
3010        #[cfg(target_arch = "aarch64")]
3011        if sdot_enabled() {
3012            let acts: [(SplitAct, SplitAct); N] =
3013                std::array::from_fn(|i| (split_act(&ctxs[i].xs1), split_act(&ctxs[i].xs2)));
3014            let closures: [_; N] = std::array::from_fn(|i| {
3015                let (c, a, o1, o2) = (&ctxs[i], &acts[i], p1[i], p2[i]);
3016                move |start: usize, end: usize| {
3017                    q8_range2_sdot(c.bytes, c.row_scale, &a.0, &a.1, c.cols, o1, o2, start, end)
3018                }
3019            });
3020            let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
3021                std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
3022            pool.run_many(&parts);
3023            return;
3024        }
3025        #[cfg(target_arch = "x86_64")]
3026        if avx2_a8w8_enabled() {
3027            let acts: [(SplitAct, SplitAct); N] =
3028                std::array::from_fn(|i| (split_act(&ctxs[i].xs1), split_act(&ctxs[i].xs2)));
3029            let closures: [_; N] = std::array::from_fn(|i| {
3030                let (c, a, o1, o2) = (&ctxs[i], &acts[i], p1[i], p2[i]);
3031                move |start: usize, end: usize| {
3032                    q8_range2_avx2(c.bytes, c.row_scale, &a.0, &a.1, c.cols, o1, o2, start, end)
3033                }
3034            });
3035            let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
3036                std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
3037            pool.run_many(&parts);
3038            return;
3039        }
3040        let closures: [_; N] = std::array::from_fn(|i| {
3041            let (c, o1, o2) = (&ctxs[i], p1[i], p2[i]);
3042            move |start: usize, end: usize| {
3043                q8_range2_f32(
3044                    c.bytes,
3045                    c.row_scale,
3046                    &c.xs1,
3047                    &c.xs2,
3048                    c.cols,
3049                    o1,
3050                    o2,
3051                    start,
3052                    end,
3053                )
3054            }
3055        });
3056        let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
3057            std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
3058        pool.run_many(&parts);
3059    }
3060
3061    /// Fused gate+up matvec with SiLU·mul: for each row r, computes
3062    /// `silu(gate·x) * (up·x)` and writes to `out[r]`. ONE pool dispatch,
3063    /// no intermediate g/u buffers, no separate silu pass. Falls back
3064    /// (returns false) for unsupported dtype combos.
3065    pub fn matvec_silu_mul(
3066        gate: &QTensor,
3067        up: &QTensor,
3068        x: &[f32],
3069        out: &mut [f32],
3070        pool: Option<&Pool>,
3071    ) -> bool {
3072        Self::matvec_silu_mul_limited(gate, up, x, out, 0.0, pool)
3073    }
3074
3075    /// Fused gate+up+SiLU with the GLM asymmetrical clamp.  `limit == 0`
3076    /// preserves the historical unclamped helper; a positive limit clamps
3077    /// `up` to both sides and `gate` only from above, matching the GLM
3078    /// SwiGLU reference.  Keeping the limit in the row kernel avoids the two
3079    /// intermediate vectors and the extra combine pass on the Q2TP experts.
3080    pub fn matvec_silu_mul_limited(
3081        gate: &QTensor,
3082        up: &QTensor,
3083        x: &[f32],
3084        out: &mut [f32],
3085        limit: f32,
3086        pool: Option<&Pool>,
3087    ) -> bool {
3088        if gate.has_prism_contract() || up.has_prism_contract() {
3089            // The fused gate/up kernels consume x directly.  Prism requires
3090            // a per-matrix signed FWHT, so the caller must use two ordinary
3091            // descriptor-aware matvecs instead of an unrotated fast path.
3092            return false;
3093        }
3094        let inter = gate.rows();
3095        debug_assert_eq!(up.rows(), inter);
3096        debug_assert_eq!(out.len(), inter);
3097        debug_assert_eq!(gate.cols(), up.cols());
3098        if !a8w8_enabled() {
3099            return false;
3100        }
3101        let act = split_act(x);
3102        let act = &act;
3103        let x_ref = x;
3104        let out_addr = SendMut(out.as_mut_ptr());
3105
3106        match (gate, up) {
3107            // Q4Block gate + Q4Block up (most common mobile q4 models)
3108            (
3109                Self::Mapped {
3110                    dtype: TensorDtype::Q4Block,
3111                    ..
3112                },
3113                Self::Mapped {
3114                    dtype: TensorDtype::Q4Block,
3115                    ..
3116                },
3117            ) => {
3118                let (gp, gs) = q4_split(gate.quant_bytes(), gate.rows(), gate.cols());
3119                let (up_p, up_s) = q4_split(up.quant_bytes(), up.rows(), up.cols());
3120                let gpr = gate.cols() / GROUP_SIZE;
3121                let cols = gate.cols();
3122                let run = move |start: usize, end: usize| {
3123                    for r in start..end {
3124                        let mut gv = dot_q4_row_i8(gp, gs, r * gpr, gpr, &act.xq) * act.sx;
3125                        let mut uv = dot_q4_row_i8(up_p, up_s, r * gpr, gpr, &act.xq) * act.sx;
3126                        for &(j, xv) in &act.outliers {
3127                            let flat = r * cols + j;
3128                            let gb = gp[flat / 2];
3129                            let gn = if flat & 1 == 0 { gb & 0x0F } else { gb >> 4 };
3130                            let gsc = f16_to_f32(u16::from_le_bytes([
3131                                gs[(flat / GROUP_SIZE) * 2],
3132                                gs[(flat / GROUP_SIZE) * 2 + 1],
3133                            ]));
3134                            gv += ((gn as i32 - 8) as f32) * gsc * xv;
3135                            let ub = up_p[flat / 2];
3136                            let un = if flat & 1 == 0 { ub & 0x0F } else { ub >> 4 };
3137                            let usc = f16_to_f32(u16::from_le_bytes([
3138                                up_s[(flat / GROUP_SIZE) * 2],
3139                                up_s[(flat / GROUP_SIZE) * 2 + 1],
3140                            ]));
3141                            uv += ((un as i32 - 8) as f32) * usc * xv;
3142                        }
3143                        let silu_g = gv / (1.0 + (-gv).exp());
3144                        // SAFETY: disjoint row ranges per worker.
3145                        unsafe { *out_addr.at(r) = silu_g * uv };
3146                    }
3147                };
3148                dispatch_rows(pool, inter, &run);
3149                true
3150            }
3151            // Q4Tiled gate + Q4Tiled up — one row pass, both tile
3152            // streams sequential, silu·mul fused (same per-row math as
3153            // `q4t_matvec`).
3154            (
3155                Self::Mapped {
3156                    dtype: TensorDtype::Q4Tiled,
3157                    ..
3158                },
3159                Self::Mapped {
3160                    dtype: TensorDtype::Q4Tiled,
3161                    ..
3162                },
3163            ) => {
3164                let g_bytes = gate.quant_bytes();
3165                let u_bytes = up.quant_bytes();
3166                let gpr = gate.cols() / GROUP_SIZE;
3167                let run = move |start: usize, end: usize| {
3168                    for r in start..end {
3169                        let mut gv = dot_q4t_row_i8(g_bytes, r, gpr, &act.xq) * act.sx;
3170                        let mut uv = dot_q4t_row_i8(u_bytes, r, gpr, &act.xq) * act.sx;
3171                        for &(j, xv) in &act.outliers {
3172                            let (w, s) = q4t_outlier(g_bytes, r, gpr, j);
3173                            gv += w * s * xv;
3174                            let (w, s) = q4t_outlier(u_bytes, r, gpr, j);
3175                            uv += w * s * xv;
3176                        }
3177                        let silu_g = gv / (1.0 + (-gv).exp());
3178                        // SAFETY: disjoint row ranges per worker.
3179                        unsafe { *out_addr.at(r) = silu_g * uv };
3180                    }
3181                };
3182                dispatch_rows(pool, inter, &run);
3183                true
3184            }
3185            // Q4TiledP gate + Q4TiledP up — the same fused row pass, with
3186            // each row's two ladders built once and spent on both streams.
3187            (
3188                Self::Mapped {
3189                    dtype: TensorDtype::Q4TiledP,
3190                    ..
3191                },
3192                Self::Mapped {
3193                    dtype: TensorDtype::Q4TiledP,
3194                    ..
3195                },
3196            ) => {
3197                let cols = gate.cols();
3198                let gpr = cols / GROUP_SIZE;
3199                let gv_view = Q4tpView::new(gate.quant_bytes(), inter, cols);
3200                let uv_view = Q4tpView::new(up.quant_bytes(), inter, cols);
3201                let run = |start: usize, end: usize| {
3202                    let (mut gsc, mut usc) = (vec![0f32; gpr], vec![0f32; gpr]);
3203                    for r in start..end {
3204                        gv_view.scales_into(r, gpr, &mut gsc);
3205                        uv_view.scales_into(r, gpr, &mut usc);
3206                        let mut gv = dot_q4tp_row_i8(gv_view.nib, r, gpr, &act.xq, &gsc) * act.sx;
3207                        let mut uv = dot_q4tp_row_i8(uv_view.nib, r, gpr, &act.xq, &usc) * act.sx;
3208                        for &(j, xv) in &act.outliers {
3209                            let (w, s) = q4tp_outlier(gv_view.nib, r, gpr, j, &gsc);
3210                            gv += w * s * xv;
3211                            let (w, s) = q4tp_outlier(uv_view.nib, r, gpr, j, &usc);
3212                            uv += w * s * xv;
3213                        }
3214                        let silu_g = gv / (1.0 + (-gv).exp());
3215                        // SAFETY: disjoint row ranges per worker.
3216                        unsafe { *out_addr.at(r) = silu_g * uv };
3217                    }
3218                };
3219                dispatch_rows(pool, inter, &run);
3220                true
3221            }
3222            // Q1 gate + Q1 up — one row pass over both sign streams,
3223            // silu·mul fused (the per-row math of `q1_range_a8w8`); the
3224            // activation group sums are shared by both streams. Without
3225            // this arm a q1 dense FFN paid two dispatches + a combine
3226            // loop — the exact barrier this function exists to remove.
3227            (
3228                Self::Mapped {
3229                    dtype: TensorDtype::Q1,
3230                    ..
3231                },
3232                Self::Mapped {
3233                    dtype: TensorDtype::Q1,
3234                    ..
3235                },
3236            ) => {
3237                let g_bytes = gate.quant_bytes();
3238                let u_bytes = up.quant_bytes();
3239                let gpr = gate.cols() / GROUP_SIZE;
3240                let gsum = q1_group_sums(&act.xq, gpr);
3241                let gsum = &gsum;
3242                let run = move |start: usize, end: usize| {
3243                    for r in start..end {
3244                        let mut gv = dot_q1_row_i8(g_bytes, r, gpr, &act.xq, gsum) * act.sx;
3245                        let mut uv = dot_q1_row_i8(u_bytes, r, gpr, &act.xq, gsum) * act.sx;
3246                        for &(j, xv) in &act.outliers {
3247                            let (w, s) = q1_outlier(g_bytes, r, gpr, j);
3248                            gv += w * s * xv;
3249                            let (w, s) = q1_outlier(u_bytes, r, gpr, j);
3250                            uv += w * s * xv;
3251                        }
3252                        let silu_g = gv / (1.0 + (-gv).exp());
3253                        // SAFETY: disjoint row ranges per worker.
3254                        unsafe { *out_addr.at(r) = silu_g * uv };
3255                    }
3256                };
3257                dispatch_rows(pool, inter, &run);
3258                true
3259            }
3260            // Q2TiledP gate + Q2TiledP up — the 2-bit expert pair (MoE
3261            // FFNs of the W2 class): one row pass, both ladders built
3262            // once, integer code dots with shared group sums.
3263            (
3264                Self::Mapped {
3265                    dtype: TensorDtype::Q2TiledP,
3266                    ..
3267                },
3268                Self::Mapped {
3269                    dtype: TensorDtype::Q2TiledP,
3270                    ..
3271                },
3272            ) => {
3273                let cols = gate.cols();
3274                let gpr = cols / GROUP_SIZE;
3275                let gv_view = Q4tpView::new_q2(gate.quant_bytes(), inter, cols);
3276                let uv_view = Q4tpView::new_q2(up.quant_bytes(), inter, cols);
3277                let gsum = q1_group_sums(&act.xq, gpr);
3278                let gsum = &gsum;
3279                let run = move |start: usize, end: usize| {
3280                    let (mut gsc, mut usc) = (vec![0f32; gpr], vec![0f32; gpr]);
3281                    for r in start..end {
3282                        gv_view.scales_into(r, gpr, &mut gsc);
3283                        uv_view.scales_into(r, gpr, &mut usc);
3284                        let mut gv =
3285                            dot_q2tp_row_i8(gv_view.nib, r, gpr, &act.xq, gsum, &gsc) * act.sx;
3286                        let mut uv =
3287                            dot_q2tp_row_i8(uv_view.nib, r, gpr, &act.xq, gsum, &usc) * act.sx;
3288                        for &(j, xv) in &act.outliers {
3289                            let (w, s) = q2tp_outlier(gv_view.nib, r, gpr, j, &gsc);
3290                            gv += w * s * xv;
3291                            let (w, s) = q2tp_outlier(uv_view.nib, r, gpr, j, &usc);
3292                            uv += w * s * xv;
3293                        }
3294                        let silu_g = gv / (1.0 + (-gv).exp());
3295                        // SAFETY: disjoint row ranges per worker.
3296                        unsafe { *out_addr.at(r) = silu_g * uv };
3297                    }
3298                };
3299                dispatch_rows(pool, inter, &run);
3300                true
3301            }
3302            // Q8Row gate + Q8Row up — one row pass over both i8 streams.
3303            // Q8_2f stays out on purpose: its column field prescales the
3304            // activations PER TENSOR, which breaks this fn's shared
3305            // split_act contract — it keeps the two-dispatch path.
3306            (
3307                Self::Mapped {
3308                    dtype: TensorDtype::Q8Row,
3309                    row_scale: g_rs,
3310                    ..
3311                },
3312                Self::Mapped {
3313                    dtype: TensorDtype::Q8Row,
3314                    row_scale: u_rs,
3315                    ..
3316                },
3317            ) => {
3318                let g_bytes = gate.quant_bytes();
3319                let u_bytes = up.quant_bytes();
3320                let cols = gate.cols();
3321                let run = move |start: usize, end: usize| {
3322                    for r in start..end {
3323                        let gv = q8_row_dot(&g_bytes[r * cols..(r + 1) * cols], act) * g_rs[r];
3324                        let uv = q8_row_dot(&u_bytes[r * cols..(r + 1) * cols], act) * u_rs[r];
3325                        let silu_g = gv / (1.0 + (-gv).exp());
3326                        // SAFETY: disjoint row ranges per worker.
3327                        unsafe { *out_addr.at(r) = silu_g * uv };
3328                    }
3329                };
3330                dispatch_rows(pool, inter, &run);
3331                true
3332            }
3333            // Q1T gate + Q1T up
3334            (
3335                Self::Mapped {
3336                    dtype: TensorDtype::Q1T,
3337                    ..
3338                },
3339                Self::Mapped {
3340                    dtype: TensorDtype::Q1T,
3341                    ..
3342                },
3343            ) => {
3344                const TILE: usize = cortiq_core::quant::Q1T_TILE;
3345                let g_bytes = gate.quant_bytes();
3346                let u_bytes = up.quant_bytes();
3347                let gpr = gate.cols() / GROUP_SIZE;
3348                let (g_rp, g_ent, g_ov) = q1t_overlay(g_bytes, inter * gpr * TILE, inter);
3349                let (u_rp, u_ent, u_ov) = q1t_overlay(u_bytes, inter * gpr * TILE, inter);
3350                let run = move |start: usize, end: usize| {
3351                    for r in start..end {
3352                        let mut gv = q1t_dot_row_i8(g_bytes, r, gpr, &act.xq) * act.sx;
3353                        let mut uv = q1t_dot_row_i8(u_bytes, r, gpr, &act.xq) * act.sx;
3354                        for &(j, xv) in &act.outliers {
3355                            gv += q1t_base_weight(g_bytes, r, gpr, j) * xv;
3356                            uv += q1t_base_weight(u_bytes, r, gpr, j) * xv;
3357                        }
3358                        gv += q1t_row_outlier_correction(g_bytes, r, g_rp, g_ent, g_ov, x_ref);
3359                        uv += q1t_row_outlier_correction(u_bytes, r, u_rp, u_ent, u_ov, x_ref);
3360                        let silu_g = gv / (1.0 + (-gv).exp());
3361                        // SAFETY: disjoint row ranges per worker.
3362                        unsafe { *out_addr.at(r) = silu_g * uv };
3363                    }
3364                };
3365                dispatch_rows(pool, inter, &run);
3366                true
3367            }
3368            _ => false,
3369        }
3370    }
3371
3372    /// Every routed expert's fused gate/up/SiLU under ONE pool dispatch.
3373    ///
3374    /// The per-expert path pays a pool barrier per expert per stage: at 9
3375    /// experts over 40 layers that is ~720 barriers a token, and a decode
3376    /// profile of Qwen3.6-35B-A3B showed the pool parked in
3377    /// `psynch_cvwait` about twice as long as it spent computing. Laying
3378    /// every expert's rows end-to-end in one virtual row space collapses
3379    /// the stage to a single dispatch. The per-row body is the
3380    /// single-expert q4tp arm verbatim, so outputs are bit-identical.
3381    ///
3382    /// `false` = something is outside the fused q4tp kernel (dtype, shape,
3383    /// or a transformed tensor); the caller walks the ordinary per-expert
3384    /// path. Float activations use the same exact scalar rows, still fused
3385    /// under one pool dispatch.
3386    pub fn moe_gate_up_many(
3387        pairs: &[(&QTensor, &QTensor)],
3388        x: &[f32],
3389        outs: &mut [Vec<f32>],
3390        pool: Option<&Pool>,
3391    ) -> bool {
3392        if pairs.is_empty() || pairs.len() != outs.len() {
3393            return false;
3394        }
3395        if !a8w8_enabled() {
3396            let groups = vec![vec![0]; pairs.len()];
3397            return Self::moe_gate_up_rows(pairs, &groups, x, outs, pool);
3398        }
3399        let inter = pairs[0].0.rows();
3400        let cols = pairs[0].0.cols();
3401        if cols % GROUP_SIZE != 0 {
3402            return false;
3403        }
3404        let gpr = cols / GROUP_SIZE;
3405        // Uniform layout across every routed pair: q4tp, or the 2-bit
3406        // profile's q2tp gate/up (the W2 class). Mixed sets refuse.
3407        let q2 = matches!(
3408            pairs[0].0,
3409            Self::Mapped {
3410                dtype: TensorDtype::Q2TiledP,
3411                ..
3412            }
3413        );
3414        let want = if q2 {
3415            TensorDtype::Q2TiledP
3416        } else {
3417            TensorDtype::Q4TiledP
3418        };
3419        let mut views = Vec::with_capacity(pairs.len() * 2);
3420        for ((g, u), o) in pairs.iter().zip(outs.iter()) {
3421            let both = matches!(g, Self::Mapped { dtype, .. } if *dtype == want)
3422                && matches!(u, Self::Mapped { dtype, .. } if *dtype == want);
3423            if !both
3424                || g.rows() != inter
3425                || u.rows() != inter
3426                || g.cols() != cols
3427                || u.cols() != cols
3428                || o.len() != inter
3429            {
3430                return false;
3431            }
3432            let mk = if q2 { Q4tpView::new_q2 } else { Q4tpView::new };
3433            views.push(mk(g.quant_bytes(), inter, cols));
3434            views.push(mk(u.quant_bytes(), inter, cols));
3435        }
3436        let act = split_act(x);
3437        let gsum = if q2 {
3438            q1_group_sums(&act.xq, gpr)
3439        } else {
3440            Vec::new()
3441        };
3442        let (act, gsum) = (&act, &gsum);
3443        let ptrs: Vec<SendMut> = outs.iter_mut().map(|o| SendMut(o.as_mut_ptr())).collect();
3444        let (views, ptrs) = (&views, &ptrs);
3445        let run = |start: usize, end: usize| {
3446            let (mut gsc, mut usc) = (vec![0f32; gpr], vec![0f32; gpr]);
3447            for flat in start..end {
3448                let (e, r) = (flat / inter, flat % inter);
3449                let gv_view = &views[e * 2];
3450                let uv_view = &views[e * 2 + 1];
3451                gv_view.scales_into(r, gpr, &mut gsc);
3452                uv_view.scales_into(r, gpr, &mut usc);
3453                let (mut gv, mut uv) = if q2 {
3454                    (
3455                        dot_q2tp_row_i8(gv_view.nib, r, gpr, &act.xq, gsum, &gsc) * act.sx,
3456                        dot_q2tp_row_i8(uv_view.nib, r, gpr, &act.xq, gsum, &usc) * act.sx,
3457                    )
3458                } else {
3459                    (
3460                        dot_q4tp_row_i8(gv_view.nib, r, gpr, &act.xq, &gsc) * act.sx,
3461                        dot_q4tp_row_i8(uv_view.nib, r, gpr, &act.xq, &usc) * act.sx,
3462                    )
3463                };
3464                for &(j, xv) in &act.outliers {
3465                    let (og, ou) = if q2 {
3466                        (
3467                            q2tp_outlier(gv_view.nib, r, gpr, j, &gsc),
3468                            q2tp_outlier(uv_view.nib, r, gpr, j, &usc),
3469                        )
3470                    } else {
3471                        (
3472                            q4tp_outlier(gv_view.nib, r, gpr, j, &gsc),
3473                            q4tp_outlier(uv_view.nib, r, gpr, j, &usc),
3474                        )
3475                    };
3476                    gv += og.0 * og.1 * xv;
3477                    uv += ou.0 * ou.1 * xv;
3478                }
3479                let silu_g = gv / (1.0 + (-gv).exp());
3480                // SAFETY: one worker owns each (expert, row) pair.
3481                unsafe { *ptrs[e].at(r) = silu_g * uv };
3482            }
3483        };
3484        dispatch_rows(pool, pairs.len() * inter, &run);
3485        true
3486    }
3487
3488    /// Every routed expert's down projection, weighted and summed into
3489    /// `out`, under ONE pool dispatch.
3490    ///
3491    /// Partitioned by OUTPUT row rather than by expert: each row is owned
3492    /// by a single worker, so the experts are summed in the caller's order
3493    /// — the same sequence of f32 adds the serial `out[i] += w·eo[i]` loop
3494    /// performs, hence bit-identical. Partitioning by expert instead would
3495    /// race on the shared accumulator.
3496    pub fn moe_down_many(
3497        downs: &[&QTensor],
3498        gs: &[Vec<f32>],
3499        weights: &[f32],
3500        out: &mut [f32],
3501        pool: Option<&Pool>,
3502    ) -> bool {
3503        if downs.is_empty() || downs.len() != gs.len() || downs.len() != weights.len() {
3504            return false;
3505        }
3506        if !a8w8_enabled() {
3507            let mut terms = vec![vec![0.0; out.len()]; downs.len()];
3508            if !Self::moe_down_rows(downs, &vec![1; downs.len()], gs, &mut terms, pool) {
3509                return false;
3510            }
3511            out.fill(0.0);
3512            for (row, &w) in terms.iter().zip(weights) {
3513                for (o, &v) in out.iter_mut().zip(row) {
3514                    *o += w * v;
3515                }
3516            }
3517            return true;
3518        }
3519        let rows = out.len();
3520        let cols = downs[0].cols();
3521        if cols % GROUP_SIZE != 0 {
3522            return false;
3523        }
3524        let gpr = cols / GROUP_SIZE;
3525        let mut views = Vec::with_capacity(downs.len());
3526        for (d, g) in downs.iter().zip(gs.iter()) {
3527            if !matches!(
3528                d,
3529                Self::Mapped {
3530                    dtype: TensorDtype::Q4TiledP,
3531                    ..
3532                }
3533            ) || d.rows() != rows
3534                || d.cols() != cols
3535                || g.len() != cols
3536            {
3537                return false;
3538            }
3539            views.push(Q4tpView::new(d.quant_bytes(), rows, cols));
3540        }
3541        // One int8 split per expert — the activation vectors differ.
3542        let acts: Vec<SplitAct> = gs.iter().map(|g| split_act(g)).collect();
3543        // Partitioned by OUTPUT row, with the experts folded inside: each
3544        // row is owned by one worker, so they are summed in the caller's
3545        // order — the same f32 sequence the serial `out[i] += w·eo[i]`
3546        // loop produces. Partitioning by expert instead would either race
3547        // on the accumulator or need a scratch plane and a second pass;
3548        // measured, that variant was a wash, so this keeps the simpler
3549        // shape.
3550        let out_addr = SendMut(out.as_mut_ptr());
3551        let (views, acts, weights) = (&views, &acts, &weights);
3552        let run = |start: usize, end: usize| {
3553            let mut sc = vec![0f32; gpr];
3554            for r in start..end {
3555                let mut acc = 0f32;
3556                for (e, v) in views.iter().enumerate() {
3557                    v.scales_into(r, gpr, &mut sc);
3558                    let a = &acts[e];
3559                    let mut d = dot_q4tp_row_i8(v.nib, r, gpr, &a.xq, &sc) * a.sx;
3560                    for &(j, xv) in &a.outliers {
3561                        let (w, s) = q4tp_outlier(v.nib, r, gpr, j, &sc);
3562                        d += w * s * xv;
3563                    }
3564                    acc += weights[e] * d;
3565                }
3566                // SAFETY: disjoint row ranges per worker.
3567                unsafe { *out_addr.at(r) = acc };
3568            }
3569        };
3570        dispatch_rows(pool, rows, &run);
3571        true
3572    }
3573
3574    /// `moe_gate_up_many` for SEVERAL tokens at once, decode-exact: expert
3575    /// `e` (`pairs[e]`, q4tp) serves the tokens `groups[e]` (row indices
3576    /// into `xs`, each `cols` wide). Every (expert, token) output is
3577    /// bit-identical to `moe_gate_up_many` run on that token alone — the
3578    /// same int8 activation split, VNNI dots, outlier terms and inline
3579    /// SiLU — while each weight row is read once for all the tokens routed
3580    /// to its expert (the speculative verify's expert sharing). `outs` is
3581    /// flat in (expert, token-of-group) order. False = not covered (not
3582    /// q4tp): the caller takes the per-token path. With float activations,
3583    /// the exact scalar row kernel replaces the int8 dot without changing
3584    /// the shared dispatch or route-order reduction.
3585    pub fn moe_gate_up_rows(
3586        pairs: &[(&QTensor, &QTensor)],
3587        groups: &[Vec<usize>],
3588        xs: &[f32],
3589        outs: &mut [Vec<f32>],
3590        pool: Option<&Pool>,
3591    ) -> bool {
3592        if pairs.is_empty() || pairs.len() != groups.len() {
3593            return false;
3594        }
3595        let inter = pairs[0].0.rows();
3596        let cols = pairs[0].0.cols();
3597        let n_pairs: usize = groups.iter().map(|g| g.len()).sum();
3598        if cols == 0 || cols % GROUP_SIZE != 0 || outs.len() != n_pairs || xs.len() % cols != 0 {
3599            return false;
3600        }
3601        let b = xs.len() / cols;
3602        let gpr = cols / GROUP_SIZE;
3603        let mut views = Vec::with_capacity(pairs.len() * 2);
3604        for (g, u) in pairs {
3605            let q4tp = |t: &QTensor| {
3606                matches!(
3607                    t,
3608                    Self::Mapped {
3609                        dtype: TensorDtype::Q4TiledP,
3610                        ..
3611                    }
3612                )
3613            };
3614            if g.has_prism_contract()
3615                || u.has_prism_contract()
3616                || !q4tp(g)
3617                || !q4tp(u)
3618                || g.rows() != inter
3619                || u.rows() != inter
3620                || g.cols() != cols
3621                || u.cols() != cols
3622            {
3623                return false;
3624            }
3625            views.push(Q4tpView::new(g.quant_bytes(), inter, cols));
3626            views.push(Q4tpView::new(u.quant_bytes(), inter, cols));
3627        }
3628        if outs.iter().any(|o| o.len() != inter) || groups.iter().flatten().any(|&t| t >= b) {
3629            return false;
3630        }
3631        let quantized = a8w8_enabled();
3632        let acts: Vec<SplitAct> = if quantized {
3633            (0..b)
3634                .map(|t| split_act(&xs[t * cols..(t + 1) * cols]))
3635                .collect()
3636        } else {
3637            Vec::new()
3638        };
3639        let mut offs = Vec::with_capacity(groups.len());
3640        let mut o = 0usize;
3641        for g in groups {
3642            offs.push(o);
3643            o += g.len();
3644        }
3645        let ptrs: Vec<SendMut> = outs.iter_mut().map(|o| SendMut(o.as_mut_ptr())).collect();
3646        let (views, ptrs, acts, offs) = (&views, &ptrs, &acts, &offs);
3647        let run = |start: usize, end: usize| {
3648            let (mut gsc, mut usc) = (vec![0f32; gpr], vec![0f32; gpr]);
3649            for flat in start..end {
3650                let (e, r) = (flat / inter, flat % inter);
3651                let (gv_view, uv_view) = (&views[e * 2], &views[e * 2 + 1]);
3652                gv_view.scales_into(r, gpr, &mut gsc);
3653                uv_view.scales_into(r, gpr, &mut usc);
3654                for (k, &t) in groups[e].iter().enumerate() {
3655                    if !quantized {
3656                        let x = &xs[t * cols..(t + 1) * cols];
3657                        let gv = q4tp_row_exact(gv_view.nib, r, gpr, x, &gsc);
3658                        let uv = q4tp_row_exact(uv_view.nib, r, gpr, x, &usc);
3659                        unsafe { *ptrs[offs[e] + k].at(r) = (gv / (1.0 + (-gv).exp())) * uv };
3660                        continue;
3661                    }
3662                    let act = &acts[t];
3663                    let mut gv = dot_q4tp_row_i8(gv_view.nib, r, gpr, &act.xq, &gsc) * act.sx;
3664                    let mut uv = dot_q4tp_row_i8(uv_view.nib, r, gpr, &act.xq, &usc) * act.sx;
3665                    for &(j, xv) in &act.outliers {
3666                        let og = q4tp_outlier(gv_view.nib, r, gpr, j, &gsc);
3667                        let ou = q4tp_outlier(uv_view.nib, r, gpr, j, &usc);
3668                        gv += og.0 * og.1 * xv;
3669                        uv += ou.0 * ou.1 * xv;
3670                    }
3671                    let silu_g = gv / (1.0 + (-gv).exp());
3672                    // SAFETY: one worker owns each (expert, row) cell of
3673                    // every output of the expert's group.
3674                    unsafe { *ptrs[offs[e] + k].at(r) = silu_g * uv };
3675                }
3676            }
3677        };
3678        dispatch_rows(pool, pairs.len() * inter, &run);
3679        true
3680    }
3681
3682    /// The per-(expert, token) down terms `moe_down_many` weights and sums,
3683    /// for SEVERAL tokens: `outs[p][o] = down_e[o] · gs[p]` (int8 split of
3684    /// `gs[p]`, VNNI dot, outlier terms — bit-identical to that kernel's
3685    /// `d`), each down row read once for its expert's whole group. The
3686    /// caller sums `w·d` per token in its route order, which reproduces
3687    /// `moe_down_many`'s f32 sequence exactly. Layout as `moe_gate_up_rows`.
3688    pub fn moe_down_rows(
3689        downs: &[&QTensor],
3690        group_lens: &[usize],
3691        gs: &[Vec<f32>],
3692        outs: &mut [Vec<f32>],
3693        pool: Option<&Pool>,
3694    ) -> bool {
3695        if downs.is_empty() || downs.len() != group_lens.len() {
3696            return false;
3697        }
3698        let rows = downs[0].rows();
3699        let cols = downs[0].cols();
3700        let n_pairs: usize = group_lens.iter().sum();
3701        if cols == 0 || cols % GROUP_SIZE != 0 || gs.len() != n_pairs || outs.len() != n_pairs {
3702            return false;
3703        }
3704        let gpr = cols / GROUP_SIZE;
3705        let mut views = Vec::with_capacity(downs.len());
3706        for d in downs {
3707            if d.has_prism_contract()
3708                || !matches!(
3709                    d,
3710                    Self::Mapped {
3711                        dtype: TensorDtype::Q4TiledP,
3712                        ..
3713                    }
3714                )
3715                || d.rows() != rows
3716                || d.cols() != cols
3717            {
3718                return false;
3719            }
3720            views.push(Q4tpView::new(d.quant_bytes(), rows, cols));
3721        }
3722        if gs.iter().any(|g| g.len() != cols) || outs.iter().any(|o| o.len() != rows) {
3723            return false;
3724        }
3725        let quantized = a8w8_enabled();
3726        let acts: Vec<SplitAct> = if quantized {
3727            gs.iter().map(|g| split_act(g)).collect()
3728        } else {
3729            Vec::new()
3730        };
3731        let mut offs = Vec::with_capacity(group_lens.len());
3732        let mut o = 0usize;
3733        for &l in group_lens {
3734            offs.push(o);
3735            o += l;
3736        }
3737        let ptrs: Vec<SendMut> = outs.iter_mut().map(|o| SendMut(o.as_mut_ptr())).collect();
3738        let (views, ptrs, acts, offs) = (&views, &ptrs, &acts, &offs);
3739        let run = |start: usize, end: usize| {
3740            let mut sc = vec![0f32; gpr];
3741            for flat in start..end {
3742                let (e, r) = (flat / rows, flat % rows);
3743                let v = &views[e];
3744                v.scales_into(r, gpr, &mut sc);
3745                for k in 0..group_lens[e] {
3746                    if !quantized {
3747                        let d = q4tp_row_exact(v.nib, r, gpr, &gs[offs[e] + k], &sc);
3748                        unsafe { *ptrs[offs[e] + k].at(r) = d };
3749                        continue;
3750                    }
3751                    let a = &acts[offs[e] + k];
3752                    let mut d = dot_q4tp_row_i8(v.nib, r, gpr, &a.xq, &sc) * a.sx;
3753                    for &(j, xv) in &a.outliers {
3754                        let (w, s) = q4tp_outlier(v.nib, r, gpr, j, &sc);
3755                        d += w * s * xv;
3756                    }
3757                    // SAFETY: one worker owns each (expert, row) cell.
3758                    unsafe { *ptrs[offs[e] + k].at(r) = d };
3759                }
3760            }
3761        };
3762        dispatch_rows(pool, downs.len() * rows, &run);
3763        true
3764    }
3765}
3766
3767/// Batched q8 kernel: same math as qmatvec, the row makes a single
3768/// pass from memory for the whole batch.
3769/// Accelerate CBLAS — the Apple AMX matrix units, the same engine
3770/// llama.cpp's `-ngl 0` prefill rides via ggml-blas.
3771#[cfg(target_os = "macos")]
3772mod accel_blas {
3773    #[link(name = "Accelerate", kind = "framework")]
3774    unsafe extern "C" {
3775        pub fn cblas_sgemm(
3776            order: i32,
3777            trans_a: i32,
3778            trans_b: i32,
3779            m: i32,
3780            n: i32,
3781            k: i32,
3782            alpha: f32,
3783            a: *const f32,
3784            lda: i32,
3785            b: *const f32,
3786            ldb: i32,
3787            beta: f32,
3788            c: *mut f32,
3789            ldc: i32,
3790        );
3791    }
3792}
3793
3794#[cfg(target_os = "macos")]
3795pub(crate) fn accel_gemm_enabled() -> bool {
3796    static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
3797    *ON.get_or_init(|| std::env::var("CMF_ACCEL").map(|v| v != "0").unwrap_or(true))
3798}
3799
3800/// Off macOS the "accel" GEMM is the portable NEON micro-kernel below —
3801/// same entry point, so the batched-attention path opens on mobile.
3802#[cfg(all(target_arch = "aarch64", not(target_os = "macos")))]
3803pub(crate) fn accel_gemm_enabled() -> bool {
3804    true
3805}
3806
3807/// Portable NEON f32 GEMM (row-major, optional Bᵀ): a 4×8 fmla
3808/// micro-kernel with A broadcast against B panels — the mobile stand-in
3809/// for Accelerate in the batched causal attention (QKᵀ and P·V). Not a
3810/// BLAS: shapes here are the attention panels (m ≤ heads·chunk,
3811/// k = head_dim or context), and the goal is removing the per-position
3812/// quadratic wall, not peak GEMM.
3813#[cfg(target_arch = "aarch64")]
3814#[allow(clippy::too_many_arguments)]
3815pub(crate) fn neon_gemm_rm(
3816    m: usize,
3817    n: usize,
3818    k: usize,
3819    alpha: f32,
3820    a: &[f32],
3821    lda: usize,
3822    b_mat: &[f32],
3823    ldb: usize,
3824    b_rows_are_n: bool,
3825    c: &mut [f32],
3826    ldc: usize,
3827) {
3828    debug_assert!(a.len() >= (m - 1) * lda + k);
3829    debug_assert!(c.len() >= (m - 1) * ldc + n);
3830    // SAFETY: bounds asserted above; NEON is baseline on aarch64.
3831    unsafe {
3832        use core::arch::aarch64::*;
3833        let mut i = 0usize;
3834        while i < m {
3835            let mi = (m - i).min(4);
3836            let mut j = 0usize;
3837            while j < n {
3838                let nj = (n - j).min(8);
3839                if mi == 4 && nj == 8 {
3840                    let (mut c0a, mut c0b) = (vdupq_n_f32(0.0), vdupq_n_f32(0.0));
3841                    let (mut c1a, mut c1b) = (vdupq_n_f32(0.0), vdupq_n_f32(0.0));
3842                    let (mut c2a, mut c2b) = (vdupq_n_f32(0.0), vdupq_n_f32(0.0));
3843                    let (mut c3a, mut c3b) = (vdupq_n_f32(0.0), vdupq_n_f32(0.0));
3844                    for p in 0..k {
3845                        let (b0, b1) = if b_rows_are_n {
3846                            // B is [n, k]: column p of Bᵀ = element p of
3847                            // eight consecutive B rows — gathered.
3848                            let base = b_mat.as_ptr().add(j * ldb + p);
3849                            let g = |o: usize| *base.add(o * ldb);
3850                            ([g(0), g(1), g(2), g(3)], [g(4), g(5), g(6), g(7)])
3851                        } else {
3852                            let base = b_mat.as_ptr().add(p * ldb + j);
3853                            (
3854                                [*base, *base.add(1), *base.add(2), *base.add(3)],
3855                                [*base.add(4), *base.add(5), *base.add(6), *base.add(7)],
3856                            )
3857                        };
3858                        let bv0 = vld1q_f32(b0.as_ptr());
3859                        let bv1 = vld1q_f32(b1.as_ptr());
3860                        let a0 = vdupq_n_f32(*a.as_ptr().add(i * lda + p));
3861                        let a1 = vdupq_n_f32(*a.as_ptr().add((i + 1) * lda + p));
3862                        let a2 = vdupq_n_f32(*a.as_ptr().add((i + 2) * lda + p));
3863                        let a3 = vdupq_n_f32(*a.as_ptr().add((i + 3) * lda + p));
3864                        c0a = vfmaq_f32(c0a, a0, bv0);
3865                        c0b = vfmaq_f32(c0b, a0, bv1);
3866                        c1a = vfmaq_f32(c1a, a1, bv0);
3867                        c1b = vfmaq_f32(c1b, a1, bv1);
3868                        c2a = vfmaq_f32(c2a, a2, bv0);
3869                        c2b = vfmaq_f32(c2b, a2, bv1);
3870                        c3a = vfmaq_f32(c3a, a3, bv0);
3871                        c3b = vfmaq_f32(c3b, a3, bv1);
3872                    }
3873                    let al = vdupq_n_f32(alpha);
3874                    for (r, (ca, cb)) in [(c0a, c0b), (c1a, c1b), (c2a, c2b), (c3a, c3b)]
3875                        .iter()
3876                        .enumerate()
3877                    {
3878                        let dst = c.as_mut_ptr().add((i + r) * ldc + j);
3879                        vst1q_f32(dst, vmulq_f32(*ca, al));
3880                        vst1q_f32(dst.add(4), vmulq_f32(*cb, al));
3881                    }
3882                } else {
3883                    for r in 0..mi {
3884                        for q in 0..nj {
3885                            let mut acc = 0f32;
3886                            for p in 0..k {
3887                                let bv = if b_rows_are_n {
3888                                    b_mat[(j + q) * ldb + p]
3889                                } else {
3890                                    b_mat[p * ldb + j + q]
3891                                };
3892                                acc += a[(i + r) * lda + p] * bv;
3893                            }
3894                            c[(i + r) * ldc + j + q] = acc * alpha;
3895                        }
3896                    }
3897                }
3898                j += nj;
3899            }
3900            i += mi;
3901        }
3902    }
3903}
3904
3905/// Off-macOS aarch64: the batched attention rides the NEON micro-GEMM.
3906#[cfg(all(target_arch = "aarch64", not(target_os = "macos")))]
3907#[allow(clippy::too_many_arguments)]
3908pub(crate) fn sgemm_rm(
3909    m: usize,
3910    n: usize,
3911    k: usize,
3912    alpha: f32,
3913    a: &[f32],
3914    lda: usize,
3915    b_mat: &[f32],
3916    ldb: usize,
3917    b_rows_are_n: bool,
3918    c: &mut [f32],
3919    ldc: usize,
3920) {
3921    neon_gemm_rm(m, n, k, alpha, a, lda, b_mat, ldb, b_rows_are_n, c, ldc);
3922}
3923
3924/// Row-major f32 GEMM, exposed for offline tools (the AWNP pass builds a
3925/// per-layer projection and applies it to every expert; a naive triple loop
3926/// would turn a two-minute job into half an hour).
3927#[allow(clippy::too_many_arguments)]
3928pub fn sgemm_public(
3929    m: usize,
3930    n: usize,
3931    k: usize,
3932    alpha: f32,
3933    a: &[f32],
3934    lda: usize,
3935    b_mat: &[f32],
3936    ldb: usize,
3937    b_rows_are_n: bool,
3938    c: &mut [f32],
3939    ldc: usize,
3940) {
3941    #[cfg(any(target_os = "macos", target_arch = "aarch64"))]
3942    {
3943        sgemm_rm(m, n, k, alpha, a, lda, b_mat, ldb, b_rows_are_n, c, ldc);
3944    }
3945    // x86 without Accelerate has no sgemm_rm: the specialized paths there are
3946    // quantized kernels, not an f32 GEMM. Only the offline AWNP pass reaches
3947    // this, so correctness matters and throughput does not — a triple loop is
3948    // the honest fallback rather than a reason to make the tool macOS-only.
3949    #[cfg(not(any(target_os = "macos", target_arch = "aarch64")))]
3950    {
3951        for i in 0..m {
3952            for j in 0..n {
3953                let mut acc = 0f32;
3954                for p in 0..k {
3955                    let bv = if b_rows_are_n {
3956                        b_mat[j * ldb + p]
3957                    } else {
3958                        b_mat[p * ldb + j]
3959                    };
3960                    acc += a[i * lda + p] * bv;
3961                }
3962                c[i * ldc + j] = alpha * acc;
3963            }
3964        }
3965    }
3966}
3967
3968/// Row-major f32 GEMM on Accelerate: C[m,n] = alpha·A[m,k] × B(ᵀ).
3969/// `b_rows_are_n` = true multiplies by Bᵀ where B is stored [n, k].
3970#[cfg(target_os = "macos")]
3971#[allow(clippy::too_many_arguments)]
3972pub(crate) fn sgemm_rm(
3973    m: usize,
3974    n: usize,
3975    k: usize,
3976    alpha: f32,
3977    a: &[f32],
3978    lda: usize,
3979    b_mat: &[f32],
3980    ldb: usize,
3981    b_rows_are_n: bool,
3982    c: &mut [f32],
3983    ldc: usize,
3984) {
3985    debug_assert!(a.len() >= (m - 1) * lda + k);
3986    debug_assert!(c.len() >= (m - 1) * ldc + n);
3987    // Test hook: route the attention GEMMs through the portable NEON
3988    // micro-kernel ON APPLE SILICON — how the mobile batched attend is
3989    // measured without a phone in the loop. (Intel macOS has no NEON —
3990    // the hook is a no-op there, Accelerate continues below.)
3991    #[cfg(target_arch = "aarch64")]
3992    if std::env::var("CMF_FORCE_NEON_GEMM")
3993        .map(|v| v == "1")
3994        .unwrap_or(false)
3995    {
3996        return neon_gemm_rm(m, n, k, alpha, a, lda, b_mat, ldb, b_rows_are_n, c, ldc);
3997    }
3998    unsafe {
3999        accel_blas::cblas_sgemm(
4000            101, // RowMajor
4001            111, // NoTrans A
4002            if b_rows_are_n { 112 } else { 111 },
4003            m as i32,
4004            n as i32,
4005            k as i32,
4006            alpha,
4007            a.as_ptr(),
4008            lda as i32,
4009            b_mat.as_ptr(),
4010            ldb as i32,
4011            0.0,
4012            c.as_mut_ptr(),
4013            ldc as i32,
4014        );
4015    }
4016}
4017
4018/// Prefill GEMM through Accelerate (macOS): dequantize q8 rows into
4019/// f32 tiles (scale folded in, pool-parallel) and multiply each tile
4020/// on the AMX with one row-major sgemm. Tiles live in cache, weights
4021/// stream once. Numerics are f32-GEMM (not the int8 dot): prefill
4022/// logits shift within f32 rounding — tolerance-class, like every
4023/// reduction-order change; decode (M=1) never takes this path.
4024#[cfg(target_os = "macos")]
4025fn qmatmat_accel(
4026    q: &[u8],
4027    row_scale: &[f32],
4028    pre: &[std::borrow::Cow<'_, [f32]>],
4029    rows: usize,
4030    cols: usize,
4031    out: &mut [f32],
4032    pool: Option<&Pool>,
4033) {
4034    // NOTE: double-buffering the dequant against the sgemm (a scoped
4035    // thread driving the pool on tile k+1 while the caller multiplies
4036    // tile k) was tried and LOST ~6%: Accelerate's sgemm is itself
4037    // multithreaded, and the dequant workers just steal its cores.
4038    const TR: usize = 2048;
4039    let b = pre.len();
4040    thread_local! {
4041        static XPANEL: std::cell::RefCell<Vec<f32>> = const { std::cell::RefCell::new(Vec::new()) };
4042        static WTILE: std::cell::RefCell<Vec<f32>> = const { std::cell::RefCell::new(Vec::new()) };
4043    }
4044    XPANEL.with(|xp| {
4045        WTILE.with(|wt| {
4046            let mut xpanel = xp.borrow_mut();
4047            xpanel.clear();
4048            for x in pre {
4049                xpanel.extend_from_slice(x);
4050            }
4051            let mut wtile = wt.borrow_mut();
4052            wtile.resize(TR * cols, 0.0);
4053            let mut r0 = 0usize;
4054            while r0 < rows {
4055                let tr = TR.min(rows - r0);
4056                // Dequant the tile (scale folded) — pool-parallel.
4057                let wt_addr = SendMut(wtile.as_mut_ptr());
4058                let run = |start: usize, end: usize| {
4059                    for r in start..end {
4060                        let row = &q[(r0 + r) * cols..(r0 + r + 1) * cols];
4061                        let s = row_scale[r0 + r];
4062                        // SAFETY: workers cover disjoint r ranges.
4063                        let dst =
4064                            unsafe { std::slice::from_raw_parts_mut(wt_addr.at(r * cols), cols) };
4065                        for (d, &v) in dst.iter_mut().zip(row) {
4066                            *d = (v as i8) as f32 * s;
4067                        }
4068                    }
4069                };
4070                dispatch_rows(pool, tr, &run);
4071                // C[b, tr] (at column r0 of out[b, rows]) = X · Wtileᵀ
4072                unsafe {
4073                    accel_blas::cblas_sgemm(
4074                        101, // RowMajor
4075                        111, // NoTrans A
4076                        112, // Trans B
4077                        b as i32,
4078                        tr as i32,
4079                        cols as i32,
4080                        1.0,
4081                        xpanel.as_ptr(),
4082                        cols as i32,
4083                        wtile.as_ptr(),
4084                        cols as i32,
4085                        0.0,
4086                        out.as_mut_ptr().add(r0),
4087                        rows as i32,
4088                    );
4089                }
4090                r0 += tr;
4091            }
4092        })
4093    });
4094}
4095
4096fn qmatmat(
4097    q: &[u8],
4098    row_scale: &[f32],
4099    pre: &[std::borrow::Cow<'_, [f32]>],
4100    rows: usize,
4101    cols: usize,
4102    out: &mut [f32],
4103    pool: Option<&Pool>,
4104) {
4105    let b = pre.len();
4106    debug_assert_eq!(out.len(), b * rows);
4107    // Big prefill batches ride the AMX (roadmap PR3): the row×batch
4108    // SDOT loop below peaks near the CPU's dot throughput, an order
4109    // below the matrix units. Small tensors and tiny test models stay
4110    // on the exact integer path.
4111    #[cfg(target_os = "macos")]
4112    if b >= 8 && rows * cols >= 500_000 && accel_gemm_enabled() {
4113        qmatmat_accel(q, row_scale, pre, rows, cols, out, pool);
4114        return;
4115    }
4116    #[cfg(target_arch = "aarch64")]
4117    if sdot_enabled() {
4118        let acts: Vec<SplitAct> = pre.iter().map(|x| split_act(x)).collect();
4119        let out_addr = SendMut(out.as_mut_ptr());
4120        // Blocked 2×4 (mobile prefill: no AMX to fall back on — this
4121        // path IS the ARM prefill GEMM off Apple silicon).
4122        let blocked_ok = blocked_enabled();
4123        let use_i8mm = i8mm_enabled();
4124        if blocked_ok {
4125            let run = |start: usize, end: usize| {
4126                let mut o = start;
4127                while o < end {
4128                    if o + 2 <= end {
4129                        let r0 = &q[o * cols..(o + 1) * cols];
4130                        let r1 = &q[(o + 1) * cols..(o + 2) * cols];
4131                        let mut bi = 0usize;
4132                        while bi + 4 <= acts.len() {
4133                            let xs = [
4134                                acts[bi].xq.as_slice(),
4135                                acts[bi + 1].xq.as_slice(),
4136                                acts[bi + 2].xq.as_slice(),
4137                                acts[bi + 3].xq.as_slice(),
4138                            ];
4139                            let d = if use_i8mm {
4140                                unsafe { dot_i8_smmla_2x4(r0, r1, xs) }
4141                            } else {
4142                                unsafe { dot_i8_sdot_2x4(r0, r1, xs) }
4143                            };
4144                            for (r, row) in [r0, r1].into_iter().enumerate() {
4145                                for k in 0..4 {
4146                                    let act = &acts[bi + k];
4147                                    let mut v = d[r][k] as f32 * act.sx;
4148                                    for &(j, xv) in &act.outliers {
4149                                        v += (row[j] as i8) as f32 * xv;
4150                                    }
4151                                    unsafe {
4152                                        *out_addr.at((bi + k) * rows + o + r) = v * row_scale[o + r]
4153                                    };
4154                                }
4155                            }
4156                            bi += 4;
4157                        }
4158                        while bi < acts.len() {
4159                            for (r, row) in [r0, r1].into_iter().enumerate() {
4160                                let v = row_dot_sdot(row, &acts[bi]) * row_scale[o + r];
4161                                unsafe { *out_addr.at(bi * rows + o + r) = v };
4162                            }
4163                            bi += 1;
4164                        }
4165                        o += 2;
4166                    } else {
4167                        let row = &q[o * cols..(o + 1) * cols];
4168                        for (bi, act) in acts.iter().enumerate() {
4169                            let v = row_dot_sdot(row, act) * row_scale[o];
4170                            unsafe { *out_addr.at(bi * rows + o) = v };
4171                        }
4172                        o += 1;
4173                    }
4174                }
4175            };
4176            dispatch_rows(pool, rows, &run);
4177            return;
4178        }
4179        let run = |start: usize, end: usize| {
4180            for o in start..end {
4181                let row = &q[o * cols..(o + 1) * cols];
4182                for (bi, act) in acts.iter().enumerate() {
4183                    let v = row_dot_sdot(row, act) * row_scale[o];
4184                    unsafe { *out_addr.at(bi * rows + o) = v };
4185                }
4186            }
4187        };
4188        dispatch_rows(pool, rows, &run);
4189        return;
4190    }
4191    // x86 A8W8 batch. Non-VNNI parts take the BLOCKED 2×4 kernel
4192    // (roadmap P0: two weight rows' abs() stay in registers across four
4193    // activation streams); VNNI machines keep the per-row bias-trick
4194    // dot, which is already throughput-bound there.
4195    #[cfg(target_arch = "x86_64")]
4196    if avx2_a8w8_enabled() {
4197        let acts: Vec<SplitAct> = pre.iter().map(|x| split_act(x)).collect();
4198        let out_addr = SendMut(out.as_mut_ptr());
4199        // CMF_X86_BLOCKED=0 forces the per-row path (paired in-process
4200        // A/B on noisy shared-vCPU hosts).
4201        let blocked_ok = blocked_enabled();
4202        if !avx512vnni_enabled() && blocked_ok && !row_exact() {
4203            let run = |start: usize, end: usize| {
4204                let mut o = start;
4205                while o < end {
4206                    if o + 2 <= end {
4207                        let r0 = &q[o * cols..(o + 1) * cols];
4208                        let r1 = &q[(o + 1) * cols..(o + 2) * cols];
4209                        let mut bi = 0usize;
4210                        while bi + 4 <= acts.len() {
4211                            let xs = [
4212                                acts[bi].xq.as_slice(),
4213                                acts[bi + 1].xq.as_slice(),
4214                                acts[bi + 2].xq.as_slice(),
4215                                acts[bi + 3].xq.as_slice(),
4216                            ];
4217                            let d = unsafe { dot_i8_i8_avx2_2x4(r0, r1, xs) };
4218                            for (r, row) in [r0, r1].into_iter().enumerate() {
4219                                for k in 0..4 {
4220                                    let act = &acts[bi + k];
4221                                    let mut v = d[r][k] as f32 * act.sx;
4222                                    for &(j, xv) in &act.outliers {
4223                                        v += (row[j] as i8) as f32 * xv;
4224                                    }
4225                                    unsafe {
4226                                        *out_addr.at((bi + k) * rows + o + r) = v * row_scale[o + r]
4227                                    };
4228                                }
4229                            }
4230                            bi += 4;
4231                        }
4232                        while bi < acts.len() {
4233                            for (r, row) in [r0, r1].into_iter().enumerate() {
4234                                let v = row_dot_avx2(row, &acts[bi]) * row_scale[o + r];
4235                                unsafe { *out_addr.at(bi * rows + o + r) = v };
4236                            }
4237                            bi += 1;
4238                        }
4239                        o += 2;
4240                    } else {
4241                        let row = &q[o * cols..(o + 1) * cols];
4242                        for (bi, act) in acts.iter().enumerate() {
4243                            let v = row_dot_avx2(row, act) * row_scale[o];
4244                            unsafe { *out_addr.at(bi * rows + o) = v };
4245                        }
4246                        o += 1;
4247                    }
4248                }
4249            };
4250            dispatch_rows(pool, rows, &run);
4251            return;
4252        }
4253        let run = |start: usize, end: usize| {
4254            for o in start..end {
4255                let row = &q[o * cols..(o + 1) * cols];
4256                for (bi, act) in acts.iter().enumerate() {
4257                    let v = row_dot_avx2(row, act) * row_scale[o];
4258                    unsafe { *out_addr.at(bi * rows + o) = v };
4259                }
4260            }
4261        };
4262        dispatch_rows(pool, rows, &run);
4263        return;
4264    }
4265    let out_addr = SendMut(out.as_mut_ptr());
4266    let run = |start: usize, end: usize| {
4267        for o in start..end {
4268            let row = &q[o * cols..(o + 1) * cols];
4269            for (bi, x) in pre.iter().enumerate() {
4270                let mut acc = 0f32;
4271                for j in 0..cols {
4272                    acc += (row[j] as i8) as f32 * x[j];
4273                }
4274                unsafe { *out_addr.at(bi * rows + o) = acc * row_scale[o] };
4275            }
4276        }
4277    };
4278    dispatch_rows(pool, rows, &run);
4279}
4280
4281/// Split rows across pool workers (shared qmatvec pattern). Self-balancing
4282/// — see `Pool::run_rows` for why a static 1/n split is wrong here.
4283fn dispatch_rows(pool: Option<&Pool>, rows: usize, run: &(dyn Fn(usize, usize) + Sync)) {
4284    match pool {
4285        Some(pool) if rows >= 256 => pool.run_rows(rows, run),
4286        _ => run(0, rows),
4287    }
4288}
4289
4290/// Split a q4_block blob into (packed nibbles, f16 group scales).
4291fn q4_split(bytes: &[u8], rows: usize, cols: usize) -> (&[u8], &[u8]) {
4292    let groups = rows * cols / GROUP_SIZE;
4293    bytes.split_at(groups * 16)
4294}
4295
4296/// SIMD unpack for the dominant vbit width B=4 (94% of rows on the
4297/// log2-shape calibration): 16 packed bytes -> 32 centered i8 values.
4298/// vbit packs MSB-first, so the HIGH nibble is the even element
4299/// (opposite of q4_block's lo-first interleave). Centering is u-7.
4300#[inline]
4301fn vbit_fill4(data: &[u8], buf: &mut [u8]) {
4302    #[cfg(target_arch = "aarch64")]
4303    unsafe {
4304        return vbit_fill4_neon(data, buf);
4305    }
4306    #[cfg(target_arch = "x86_64")]
4307    if avx2_enabled() {
4308        return unsafe { vbit_fill4_avx2(data, buf) };
4309    }
4310    #[allow(unreachable_code)]
4311    for (blk, chunk) in buf.chunks_exact_mut(8).enumerate() {
4312        let u = unpack8::<4>(&data[blk * 4..]);
4313        for k in 0..8 {
4314            chunk[k] = (u[k] - 7) as i8 as u8;
4315        }
4316    }
4317}
4318
4319#[cfg(target_arch = "aarch64")]
4320#[target_feature(enable = "neon")]
4321unsafe fn vbit_fill4_neon(data: &[u8], buf: &mut [u8]) {
4322    // SAFETY: buf.len() is a multiple of GROUP_SIZE=32; data holds
4323    // buf.len()/2 packed bytes (validated at load).
4324    unsafe {
4325        use core::arch::aarch64::*;
4326        let n = buf.len();
4327        let mask = vdupq_n_u8(0x0F);
4328        let seven = vdupq_n_s8(7);
4329        let mut g = 0usize;
4330        while g * 32 + 32 <= n {
4331            let b = vld1q_u8(data.as_ptr().add(g * 16));
4332            let hi = vshrq_n_u8::<4>(b);
4333            let lo = vandq_u8(b, mask);
4334            let z0 = vsubq_s8(vreinterpretq_s8_u8(vzip1q_u8(hi, lo)), seven);
4335            let z1 = vsubq_s8(vreinterpretq_s8_u8(vzip2q_u8(hi, lo)), seven);
4336            vst1q_u8(buf.as_mut_ptr().add(g * 32), vreinterpretq_u8_s8(z0));
4337            vst1q_u8(buf.as_mut_ptr().add(g * 32 + 16), vreinterpretq_u8_s8(z1));
4338            g += 1;
4339        }
4340    }
4341}
4342
4343#[cfg(target_arch = "x86_64")]
4344#[target_feature(enable = "avx2")]
4345unsafe fn vbit_fill4_avx2(data: &[u8], buf: &mut [u8]) {
4346    // SAFETY: see vbit_fill4_neon.
4347    unsafe {
4348        use core::arch::x86_64::*;
4349        let n = buf.len();
4350        let mask = _mm_set1_epi8(0x0F);
4351        let seven = _mm256_set1_epi8(7);
4352        let mut g = 0usize;
4353        while g * 32 + 32 <= n {
4354            let b = _mm_loadu_si128(data.as_ptr().add(g * 16) as *const __m128i);
4355            let hi = _mm_and_si128(_mm_srli_epi16::<4>(b), mask);
4356            let lo = _mm_and_si128(b, mask);
4357            let z = _mm256_sub_epi8(
4358                _mm256_set_m128i(_mm_unpackhi_epi8(hi, lo), _mm_unpacklo_epi8(hi, lo)),
4359                seven,
4360            );
4361            _mm256_storeu_si256(buf.as_mut_ptr().add(g * 32) as *mut __m256i, z);
4362            g += 1;
4363        }
4364    }
4365}
4366
4367/// Unpack 8 MSB-first B-bit values from exactly B bytes (fixed shifts —
4368/// no serial bit-buffer, auto-vectorizable). Every 32-value group starts
4369/// byte-aligned (32·B/8 is integral for B∈3..8), so groups decompose
4370/// into 4 such blocks.
4371#[inline(always)]
4372fn unpack8<const B: usize>(data: &[u8]) -> [i32; 8] {
4373    let mut acc = 0u64;
4374    for i in 0..B {
4375        acc = (acc << 8) | data[i] as u64;
4376    }
4377    let mask = (1u64 << B) - 1;
4378    let mut out = [0i32; 8];
4379    for (k, o) in out.iter_mut().enumerate() {
4380        *o = ((acc >> ((7 - k) * B)) & mask) as i32;
4381    }
4382    out
4383}
4384
4385/// Fused vbit matvec straight from the mapped bytes (spec §3, P13
4386/// FIG.3): [u8 bits: rows][f16 scales: rows·cols/32][bit-packed rows,
4387/// MSB-first, byte-padded]. Row data offsets are precomputed at load
4388/// (`vbit_row_offsets`) — the per-call prefix scan was O(rows) pure
4389/// overhead on every matvec.
4390#[allow(clippy::too_many_arguments)]
4391fn vbitmatvec(
4392    bytes: &[u8],
4393    offsets: &[usize],
4394    x: &[f32],
4395    rows: usize,
4396    cols: usize,
4397    out: &mut [f32],
4398    pool: Option<&Pool>,
4399) {
4400    debug_assert_eq!(out.len(), rows);
4401    debug_assert_eq!(offsets.len(), rows + 1);
4402
4403    // SDOT path: unpack the row to centered i8 once, then per-group
4404    // int8 dot against the quantized activations — same A8W8 contract
4405    // as q8 (bounded noise; CMF_SDOT=0 keeps the exact scalar path).
4406    if a8w8_enabled() {
4407        let act = split_act(x);
4408        let out_addr = SendMut(out.as_mut_ptr());
4409        let run = move |start: usize, end: usize| {
4410            vbit_range_a8w8(bytes, offsets, x, &act, rows, cols, out_addr, start, end)
4411        };
4412        dispatch_rows(pool, rows, &run);
4413        return;
4414    }
4415
4416    let out_addr = SendMut(out.as_mut_ptr());
4417    let run = move |start: usize, end: usize| {
4418        vbit_range_f32(bytes, offsets, x, rows, cols, out_addr, start, end)
4419    };
4420    dispatch_rows(pool, rows, &run);
4421}
4422
4423/// One vbit row range via the A8W8 int8 path — kernel body of
4424/// `vbitmatvec`, extracted so multi-matrix jobs can drive it for
4425/// several tensors in one dispatch (b=8 rows go exact f32).
4426#[allow(clippy::too_many_arguments)]
4427fn vbit_range_a8w8(
4428    bytes: &[u8],
4429    offsets: &[usize],
4430    x: &[f32],
4431    act: &SplitAct,
4432    rows: usize,
4433    cols: usize,
4434    out: SendMut,
4435    start: usize,
4436    end: usize,
4437) {
4438    let ng = cols / GROUP_SIZE;
4439    let bits = &bytes[..rows];
4440    let sc_off = rows;
4441    let row_dot = |r: usize| -> f32 {
4442        let b = bits[r] as usize;
4443        let l = (1i32 << (b - 1)) - 1;
4444        let mask = (1u64 << b) - 1;
4445        let data = &bytes[offsets[r]..offsets[r + 1]];
4446        if b == 8 {
4447            // u−L reaches 128 → does not fit i8; exact f32 path.
4448            let (mut acc, mut nbits, mut idx) = (0u64, 0usize, 0usize);
4449            let mut dot = 0f32;
4450            for g in 0..ng {
4451                let so = (r * ng + g) * 2;
4452                let sgf = f16_to_f32(u16::from_le_bytes([
4453                    bytes[sc_off + so],
4454                    bytes[sc_off + so + 1],
4455                ]));
4456                let xg = &x[g * GROUP_SIZE..(g + 1) * GROUP_SIZE];
4457                let mut gd = 0f32;
4458                for &xv in xg.iter() {
4459                    if nbits < 8 {
4460                        acc = (acc << 8) | data[idx] as u64;
4461                        idx += 1;
4462                        nbits += 8;
4463                    }
4464                    let u = ((acc >> (nbits - 8)) & 0xFF) as i32;
4465                    nbits -= 8;
4466                    gd += (u - l) as f32 * xv;
4467                }
4468                dot += gd * sgf;
4469            }
4470            return dot;
4471        }
4472        // Per-worker scratch: this closure runs for every row of the
4473        // tensor (lm_head ≈ 150k rows/token) — a heap allocation per
4474        // row was measurable pure overhead.
4475        thread_local! {
4476            static VBIT_SCRATCH: std::cell::RefCell<Vec<u8>> =
4477                const { std::cell::RefCell::new(Vec::new()) };
4478        }
4479        #[inline(always)]
4480        fn fill<const B: usize>(data: &[u8], l: i32, buf: &mut [u8]) {
4481            for (blk, chunk) in buf.chunks_exact_mut(8).enumerate() {
4482                let u = unpack8::<B>(&data[blk * B..]);
4483                for k in 0..8 {
4484                    chunk[k] = (u[k] - l) as i8 as u8;
4485                }
4486            }
4487        }
4488        let _ = mask;
4489        VBIT_SCRATCH.with(|scratch| {
4490            let mut buf = scratch.borrow_mut();
4491            buf.resize(cols, 0);
4492            match b {
4493                3 => fill::<3>(data, l, &mut buf),
4494                4 => vbit_fill4(data, &mut buf),
4495                5 => fill::<5>(data, l, &mut buf),
4496                6 => fill::<6>(data, l, &mut buf),
4497                _ => unreachable!(),
4498            }
4499            let mut dot = 0f32;
4500            for g in 0..ng {
4501                let so = (r * ng + g) * 2;
4502                let s = f16_to_f32(u16::from_le_bytes([
4503                    bytes[sc_off + so],
4504                    bytes[sc_off + so + 1],
4505                ]));
4506                let d = dot_i8_i8(
4507                    &buf[g * GROUP_SIZE..(g + 1) * GROUP_SIZE],
4508                    &act.xq[g * GROUP_SIZE..(g + 1) * GROUP_SIZE],
4509                ) as f32
4510                    * act.sx;
4511                dot += d * s;
4512            }
4513            for &(j, xv) in &act.outliers {
4514                let so = (r * ng + j / GROUP_SIZE) * 2;
4515                let s = f16_to_f32(u16::from_le_bytes([
4516                    bytes[sc_off + so],
4517                    bytes[sc_off + so + 1],
4518                ]));
4519                // xq is zeroed at outlier slots — add the exact term.
4520                dot += (buf[j] as i8) as f32 * s * xv;
4521            }
4522            dot
4523        })
4524    };
4525    for r in start..end {
4526        // SAFETY: disjoint row ranges per worker.
4527        unsafe { *out.at(r) = row_dot(r) };
4528    }
4529}
4530
4531/// Exact scalar vbit row range (same extraction, non-SDOT path).
4532#[allow(clippy::too_many_arguments)]
4533fn vbit_range_f32(
4534    bytes: &[u8],
4535    offsets: &[usize],
4536    x: &[f32],
4537    rows: usize,
4538    cols: usize,
4539    out: SendMut,
4540    start: usize,
4541    end: usize,
4542) {
4543    let ng = cols / GROUP_SIZE;
4544    let bits = &bytes[..rows];
4545    let sc_off = rows;
4546    // Per-bit-width specialized inner loops: the compiler unrolls the
4547    // constant shifts (the generic bit-buffer loop was branch-bound —
4548    // 5.6 vs 13.2 tok/s q4 on the 0.8B).
4549    #[inline(always)]
4550    fn dot_row<const B: usize>(
4551        data: &[u8],
4552        bytes: &[u8],
4553        sc_off: usize,
4554        r: usize,
4555        ng: usize,
4556        x: &[f32],
4557    ) -> f32 {
4558        let l = ((1i32 << (B - 1)) - 1) as f32;
4559        let gbytes = GROUP_SIZE * B / 8;
4560        let mut dot = 0f32;
4561        for g in 0..ng {
4562            let so = (r * ng + g) * 2;
4563            let s = f16_to_f32(u16::from_le_bytes([
4564                bytes[sc_off + so],
4565                bytes[sc_off + so + 1],
4566            ]));
4567            let xg = &x[g * GROUP_SIZE..(g + 1) * GROUP_SIZE];
4568            let gd0 = &data[g * gbytes..(g + 1) * gbytes];
4569            let mut gd = 0f32;
4570            for blk in 0..GROUP_SIZE / 8 {
4571                let u = unpack8::<B>(&gd0[blk * B..]);
4572                let xb = &xg[blk * 8..blk * 8 + 8];
4573                for k in 0..8 {
4574                    gd += (u[k] as f32 - l) * xb[k];
4575                }
4576            }
4577            dot += gd * s;
4578        }
4579        dot
4580    }
4581    for r in start..end {
4582        let data = &bytes[offsets[r]..offsets[r + 1]];
4583        let v = match bits[r] {
4584            3 => dot_row::<3>(data, bytes, sc_off, r, ng, x),
4585            4 => dot_row::<4>(data, bytes, sc_off, r, ng, x),
4586            5 => dot_row::<5>(data, bytes, sc_off, r, ng, x),
4587            6 => dot_row::<6>(data, bytes, sc_off, r, ng, x),
4588            8 => dot_row::<8>(data, bytes, sc_off, r, ng, x),
4589            b => unreachable!("vbit bit-width {b} (validated at load)"),
4590        };
4591        // SAFETY: disjoint row ranges per worker.
4592        unsafe { *out.at(r) = v };
4593    }
4594}
4595
4596/// Fused two-input vbit matvec: each row is unpacked from the mmap ONCE
4597/// and dotted against BOTH activations (MTP verify / pair prefill used
4598/// to run two full matvecs — double weight traffic and double unpack).
4599/// Per-input math is identical to `vbitmatvec` → same accuracy contract.
4600#[allow(clippy::too_many_arguments)]
4601fn vbitmatvec2(
4602    bytes: &[u8],
4603    offsets: &[usize],
4604    x1: &[f32],
4605    x2: &[f32],
4606    rows: usize,
4607    cols: usize,
4608    o1: &mut [f32],
4609    o2: &mut [f32],
4610    pool: Option<&Pool>,
4611) {
4612    debug_assert_eq!(o1.len(), rows);
4613    debug_assert_eq!(o2.len(), rows);
4614
4615    if a8w8_enabled() {
4616        let a1 = split_act(x1);
4617        let a2 = split_act(x2);
4618        let p1 = SendMut(o1.as_mut_ptr());
4619        let p2 = SendMut(o2.as_mut_ptr());
4620        let run = move |start: usize, end: usize| {
4621            vbit_range2_a8w8(
4622                bytes, offsets, x1, x2, &a1, &a2, rows, cols, p1, p2, start, end,
4623            )
4624        };
4625        dispatch_rows(pool, rows, &run);
4626        return;
4627    }
4628
4629    let p1 = SendMut(o1.as_mut_ptr());
4630    let p2 = SendMut(o2.as_mut_ptr());
4631    let run = move |start: usize, end: usize| {
4632        vbit_range2_f32(bytes, offsets, x1, x2, rows, cols, p1, p2, start, end)
4633    };
4634    dispatch_rows(pool, rows, &run);
4635}
4636
4637/// Two-input vbit row range via the A8W8 int8 path — kernel body of
4638/// `vbitmatvec2`, extracted for pair multi-matrix jobs (b=8 rows go
4639/// exact f32 for both lanes, bits streamed once).
4640#[allow(clippy::too_many_arguments)]
4641fn vbit_range2_a8w8(
4642    bytes: &[u8],
4643    offsets: &[usize],
4644    x1: &[f32],
4645    x2: &[f32],
4646    a1: &SplitAct,
4647    a2: &SplitAct,
4648    rows: usize,
4649    cols: usize,
4650    p1: SendMut,
4651    p2: SendMut,
4652    start: usize,
4653    end: usize,
4654) {
4655    let ng = cols / GROUP_SIZE;
4656    let bits = &bytes[..rows];
4657    let sc_off = rows;
4658    let row_dots = |r: usize| -> (f32, f32) {
4659        let b = bits[r] as usize;
4660        let l = (1i32 << (b - 1)) - 1;
4661        let data = &bytes[offsets[r]..offsets[r + 1]];
4662        if b == 8 {
4663            // u−L reaches 128 → does not fit i8; exact f32 path,
4664            // bits still streamed once for both lanes.
4665            let (mut acc, mut nbits, mut idx) = (0u64, 0usize, 0usize);
4666            let (mut d1, mut d2) = (0f32, 0f32);
4667            for g in 0..ng {
4668                let so = (r * ng + g) * 2;
4669                let sgf = f16_to_f32(u16::from_le_bytes([
4670                    bytes[sc_off + so],
4671                    bytes[sc_off + so + 1],
4672                ]));
4673                let (mut g1, mut g2) = (0f32, 0f32);
4674                for k in 0..GROUP_SIZE {
4675                    if nbits < 8 {
4676                        acc = (acc << 8) | data[idx] as u64;
4677                        idx += 1;
4678                        nbits += 8;
4679                    }
4680                    let u = ((acc >> (nbits - 8)) & 0xFF) as i32;
4681                    nbits -= 8;
4682                    let w = (u - l) as f32;
4683                    g1 += w * x1[g * GROUP_SIZE + k];
4684                    g2 += w * x2[g * GROUP_SIZE + k];
4685                }
4686                d1 += g1 * sgf;
4687                d2 += g2 * sgf;
4688            }
4689            return (d1, d2);
4690        }
4691        thread_local! {
4692            static VBIT_SCRATCH2: std::cell::RefCell<Vec<u8>> =
4693                const { std::cell::RefCell::new(Vec::new()) };
4694        }
4695        #[inline(always)]
4696        fn fill<const B: usize>(data: &[u8], l: i32, buf: &mut [u8]) {
4697            for (blk, chunk) in buf.chunks_exact_mut(8).enumerate() {
4698                let u = unpack8::<B>(&data[blk * B..]);
4699                for k in 0..8 {
4700                    chunk[k] = (u[k] - l) as i8 as u8;
4701                }
4702            }
4703        }
4704        VBIT_SCRATCH2.with(|scratch| {
4705            let mut buf = scratch.borrow_mut();
4706            buf.resize(cols, 0);
4707            match b {
4708                3 => fill::<3>(data, l, &mut buf),
4709                4 => vbit_fill4(data, &mut buf),
4710                5 => fill::<5>(data, l, &mut buf),
4711                6 => fill::<6>(data, l, &mut buf),
4712                _ => unreachable!(),
4713            }
4714            let (mut d1, mut d2) = (0f32, 0f32);
4715            for g in 0..ng {
4716                let so = (r * ng + g) * 2;
4717                let s = f16_to_f32(u16::from_le_bytes([
4718                    bytes[sc_off + so],
4719                    bytes[sc_off + so + 1],
4720                ]));
4721                let wg = &buf[g * GROUP_SIZE..(g + 1) * GROUP_SIZE];
4722                let v1 = dot_i8_i8(wg, &a1.xq[g * GROUP_SIZE..(g + 1) * GROUP_SIZE]) as f32 * a1.sx;
4723                let v2 = dot_i8_i8(wg, &a2.xq[g * GROUP_SIZE..(g + 1) * GROUP_SIZE]) as f32 * a2.sx;
4724                d1 += v1 * s;
4725                d2 += v2 * s;
4726            }
4727            for &(j, xv) in &a1.outliers {
4728                let so = (r * ng + j / GROUP_SIZE) * 2;
4729                let s = f16_to_f32(u16::from_le_bytes([
4730                    bytes[sc_off + so],
4731                    bytes[sc_off + so + 1],
4732                ]));
4733                d1 += (buf[j] as i8) as f32 * s * xv;
4734            }
4735            for &(j, xv) in &a2.outliers {
4736                let so = (r * ng + j / GROUP_SIZE) * 2;
4737                let s = f16_to_f32(u16::from_le_bytes([
4738                    bytes[sc_off + so],
4739                    bytes[sc_off + so + 1],
4740                ]));
4741                d2 += (buf[j] as i8) as f32 * s * xv;
4742            }
4743            (d1, d2)
4744        })
4745    };
4746    for r in start..end {
4747        let (v1, v2) = row_dots(r);
4748        // SAFETY: disjoint row ranges per worker.
4749        unsafe {
4750            *p1.at(r) = v1;
4751            *p2.at(r) = v2;
4752        }
4753    }
4754}
4755
4756/// Two-input exact scalar vbit row range (same extraction) —
4757/// per-bit-width specialized, two accumulators per row; per-lane
4758/// accumulation order matches `vbitmatvec` exactly.
4759#[allow(clippy::too_many_arguments)]
4760fn vbit_range2_f32(
4761    bytes: &[u8],
4762    offsets: &[usize],
4763    x1: &[f32],
4764    x2: &[f32],
4765    rows: usize,
4766    cols: usize,
4767    p1: SendMut,
4768    p2: SendMut,
4769    start: usize,
4770    end: usize,
4771) {
4772    let ng = cols / GROUP_SIZE;
4773    let bits = &bytes[..rows];
4774    let sc_off = rows;
4775    #[inline(always)]
4776    #[allow(clippy::too_many_arguments)]
4777    fn dot_row2<const B: usize>(
4778        data: &[u8],
4779        bytes: &[u8],
4780        sc_off: usize,
4781        r: usize,
4782        ng: usize,
4783        x1: &[f32],
4784        x2: &[f32],
4785    ) -> (f32, f32) {
4786        let l = ((1i32 << (B - 1)) - 1) as f32;
4787        let gbytes = GROUP_SIZE * B / 8;
4788        let (mut d1, mut d2) = (0f32, 0f32);
4789        for g in 0..ng {
4790            let so = (r * ng + g) * 2;
4791            let s = f16_to_f32(u16::from_le_bytes([
4792                bytes[sc_off + so],
4793                bytes[sc_off + so + 1],
4794            ]));
4795            let x1g = &x1[g * GROUP_SIZE..(g + 1) * GROUP_SIZE];
4796            let x2g = &x2[g * GROUP_SIZE..(g + 1) * GROUP_SIZE];
4797            let gd0 = &data[g * gbytes..(g + 1) * gbytes];
4798            let (mut g1, mut g2) = (0f32, 0f32);
4799            for blk in 0..GROUP_SIZE / 8 {
4800                let u = unpack8::<B>(&gd0[blk * B..]);
4801                for k in 0..8 {
4802                    let w = u[k] as f32 - l;
4803                    g1 += w * x1g[blk * 8 + k];
4804                    g2 += w * x2g[blk * 8 + k];
4805                }
4806            }
4807            d1 += g1 * s;
4808            d2 += g2 * s;
4809        }
4810        (d1, d2)
4811    }
4812    for r in start..end {
4813        let data = &bytes[offsets[r]..offsets[r + 1]];
4814        let (v1, v2) = match bits[r] {
4815            3 => dot_row2::<3>(data, bytes, sc_off, r, ng, x1, x2),
4816            4 => dot_row2::<4>(data, bytes, sc_off, r, ng, x1, x2),
4817            5 => dot_row2::<5>(data, bytes, sc_off, r, ng, x1, x2),
4818            6 => dot_row2::<6>(data, bytes, sc_off, r, ng, x1, x2),
4819            8 => dot_row2::<8>(data, bytes, sc_off, r, ng, x1, x2),
4820            b => unreachable!("vbit bit-width {b} (validated at load)"),
4821        };
4822        // SAFETY: disjoint row ranges per worker.
4823        unsafe {
4824            *p1.at(r) = v1;
4825            *p2.at(r) = v2;
4826        }
4827    }
4828}
4829
4830// ───────────────────── q4_tiled kernels (§4.3) ─────────────────────
4831
4832/// One q4_tiled row dot on the A8W8 int8 path: per 32-group the tile
4833/// is ONE sequential read — [f16 scale][16B nibbles] — versus the two
4834/// distant streams of the split layout. Values/order identical to the
4835/// split kernels.
4836#[inline]
4837#[allow(unreachable_code)]
4838fn dot_q4t_row_i8(bytes: &[u8], r: usize, gpr: usize, xq: &[i8]) -> f32 {
4839    #[cfg(target_arch = "aarch64")]
4840    unsafe {
4841        return dot_q4t_row_sdot(bytes, r, gpr, xq);
4842    }
4843    #[cfg(target_arch = "x86_64")]
4844    unsafe {
4845        if vnni_tiles_enabled() {
4846            return dot_q4t_row_vnni(bytes, r, gpr, xq);
4847        }
4848        return dot_q4t_row_avx2(bytes, r, gpr, xq);
4849    }
4850    let mut acc = 0f32;
4851    for gi in 0..gpr {
4852        let tile = &bytes[(r * gpr + gi) * Q4_TILE..(r * gpr + gi + 1) * Q4_TILE];
4853        let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
4854        let mut d = 0i32;
4855        for (k, &b) in tile[2..].iter().enumerate() {
4856            d += ((b & 0x0F) as i32 - 8) * xq[gi * GROUP_SIZE + k * 2] as i32
4857                + (((b >> 4) & 0x0F) as i32 - 8) * xq[gi * GROUP_SIZE + k * 2 + 1] as i32;
4858        }
4859        acc += d as f32 * s;
4860    }
4861    acc
4862}
4863
4864#[cfg(target_arch = "aarch64")]
4865#[target_feature(enable = "neon,dotprod")]
4866unsafe fn dot_q4t_row_sdot(bytes: &[u8], r: usize, gpr: usize, xq: &[i8]) -> f32 {
4867    // SAFETY: callers uphold slice-length contracts (18B tile per group,
4868    // xq.len() == gpr·GROUP_SIZE).
4869    unsafe {
4870        use core::arch::aarch64::*;
4871        use core::arch::asm;
4872        let lomask = vdupq_n_u8(0x0F);
4873        let eight = vdupq_n_s8(8);
4874        let mut acc = 0f32;
4875        for gi in 0..gpr {
4876            let t = bytes.as_ptr().add((r * gpr + gi) * Q4_TILE);
4877            let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
4878            let b = vld1q_u8(t.add(2));
4879            let lo = vandq_u8(b, lomask);
4880            let hi = vshrq_n_u8::<4>(b);
4881            let e0 = vsubq_s8(vreinterpretq_s8_u8(vzip1q_u8(lo, hi)), eight);
4882            let e1 = vsubq_s8(vreinterpretq_s8_u8(vzip2q_u8(lo, hi)), eight);
4883            let x0 = vld1q_s8(xq.as_ptr().add(gi * GROUP_SIZE));
4884            let x1 = vld1q_s8(xq.as_ptr().add(gi * GROUP_SIZE + 16));
4885            let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
4886            asm!(
4887                "sdot {a0:v}.4s, {e0:v}.16b, {x0:v}.16b",
4888                "sdot {a1:v}.4s, {e1:v}.16b, {x1:v}.16b",
4889                a0 = inout(vreg) a0, a1 = inout(vreg) a1,
4890                e0 = in(vreg) e0, x0 = in(vreg) x0, e1 = in(vreg) e1, x1 = in(vreg) x1,
4891                options(pure, nomem, nostack),
4892            );
4893            acc += vaddvq_s32(vaddq_s32(a0, a1)) as f32 * s;
4894        }
4895        acc
4896    }
4897}
4898
4899#[cfg(target_arch = "x86_64")]
4900#[target_feature(enable = "avx2")]
4901unsafe fn dot_q4t_row_avx2(bytes: &[u8], r: usize, gpr: usize, xq: &[i8]) -> f32 {
4902    // SAFETY: see dot_q4t_row_sdot.
4903    unsafe {
4904        use core::arch::x86_64::*;
4905        let lomask = _mm_set1_epi8(0x0F);
4906        let eight = _mm256_set1_epi8(8);
4907        let ones = _mm256_set1_epi16(1);
4908        let mut acc = 0f32;
4909        for gi in 0..gpr {
4910            let t = bytes.as_ptr().add((r * gpr + gi) * Q4_TILE);
4911            let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
4912            let b = _mm_loadu_si128(t.add(2) as *const __m128i);
4913            let lo = _mm_and_si128(b, lomask);
4914            let hi = _mm_and_si128(_mm_srli_epi16::<4>(b), lomask);
4915            let w = _mm256_sub_epi8(
4916                _mm256_set_m128i(_mm_unpackhi_epi8(lo, hi), _mm_unpacklo_epi8(lo, hi)),
4917                eight,
4918            );
4919            let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
4920            let p16 = _mm256_maddubs_epi16(_mm256_abs_epi8(w), _mm256_sign_epi8(x, w));
4921            let d = _mm256_madd_epi16(p16, ones);
4922            let hi128 = _mm256_extracti128_si256::<1>(d);
4923            let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
4924            let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
4925            let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
4926            acc += _mm_cvtsi128_si32(s32) as f32 * s;
4927        }
4928        acc
4929    }
4930}
4931
4932/// VNNI twin of `dot_q4t_row_avx2`: same unpack, `vpdpbusd` replaces
4933/// the maddubs+madd pair (see `dpbusd_hsum` — sums are bit-identical).
4934/// 256-bit VL encoding, so the VEX `vpsignb` stays usable.
4935#[cfg(target_arch = "x86_64")]
4936#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
4937unsafe fn dot_q4t_row_vnni(bytes: &[u8], r: usize, gpr: usize, xq: &[i8]) -> f32 {
4938    // SAFETY: see dot_q4t_row_sdot.
4939    unsafe {
4940        use core::arch::x86_64::*;
4941        let lomask = _mm_set1_epi8(0x0F);
4942        let eight = _mm256_set1_epi8(8);
4943        let mut acc = 0f32;
4944        for gi in 0..gpr {
4945            let t = bytes.as_ptr().add((r * gpr + gi) * Q4_TILE);
4946            let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
4947            let b = _mm_loadu_si128(t.add(2) as *const __m128i);
4948            let lo = _mm_and_si128(b, lomask);
4949            let hi = _mm_and_si128(_mm_srli_epi16::<4>(b), lomask);
4950            let w = _mm256_sub_epi8(
4951                _mm256_set_m128i(_mm_unpackhi_epi8(lo, hi), _mm_unpacklo_epi8(lo, hi)),
4952                eight,
4953            );
4954            let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
4955            let d = dpbusd_hsum(_mm256_abs_epi8(w), _mm256_sign_epi8(x, w));
4956            acc += d as f32 * s;
4957        }
4958        acc
4959    }
4960}
4961
4962/// One q4_tiled row against FOUR activation streams: the nibble unpack
4963/// and abs() happen once per group instead of once per (group,
4964/// activation) — the unpack is the dominant per-element cost of the
4965/// tiled format (roadmap P0 portable blocking, q4t leg).
4966#[cfg(target_arch = "x86_64")]
4967// `fma` is NOT implied by `avx2`: without it LLVM lowers _mm256_fmadd_ps
4968// to a libm call per lane — measured 2x slower than the reduction this
4969// kernel replaces. The runtime gate (`avx2_enabled`) already requires
4970// both features, so declaring it here is safe.
4971#[target_feature(enable = "avx2,fma")]
4972unsafe fn dot_q4t_row_1x4_avx2(bytes: &[u8], r: usize, gpr: usize, xs: [&[i8]; 4]) -> [f32; 4] {
4973    // SAFETY: callers uphold the 18B-tile and xq-length contracts.
4974    unsafe {
4975        use core::arch::x86_64::*;
4976        let lomask = _mm_set1_epi8(0x0F);
4977        let eight = _mm256_set1_epi8(8);
4978        let ones = _mm256_set1_epi16(1);
4979        // One f32 accumulator VECTOR per activation, reduced once at the
4980        // end. Folding each group's i32 lanes to a scalar inside the loop
4981        // costs an extracti128 + three shift/add + a movd — a cross-lane
4982        // dependency chain per (group, activation), 288 of them per row at
4983        // cols=2304. The per-group scale is what forces a float
4984        // accumulator; it does not force a horizontal sum.
4985        //
4986        // The four accumulators are NAMED, not an array: as `[__m256; 4]`
4987        // indexed by a loop variable LLVM keeps them in memory and every
4988        // group pays four 32-byte loads and stores. That alone made this
4989        // kernel 2x SLOWER than the per-group reduction it replaces
4990        // (measured on the EPYC box: 150 s vs 71 s for two 256² steps).
4991        let mut f0 = _mm256_setzero_ps();
4992        let mut f1 = _mm256_setzero_ps();
4993        let mut f2 = _mm256_setzero_ps();
4994        let mut f3 = _mm256_setzero_ps();
4995        for gi in 0..gpr {
4996            let t = bytes.as_ptr().add((r * gpr + gi) * Q4_TILE);
4997            let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
4998            let sv = _mm256_set1_ps(s);
4999            let bb = _mm_loadu_si128(t.add(2) as *const __m128i);
5000            let lo = _mm_and_si128(bb, lomask);
5001            let hi = _mm_and_si128(_mm_srli_epi16::<4>(bb), lomask);
5002            let w = _mm256_sub_epi8(
5003                _mm256_set_m128i(_mm_unpackhi_epi8(lo, hi), _mm_unpacklo_epi8(lo, hi)),
5004                eight,
5005            );
5006            let aw = _mm256_abs_epi8(w);
5007            let off = gi * GROUP_SIZE;
5008            let dot = |xq: &[i8]| {
5009                let x = _mm256_loadu_si256(xq.as_ptr().add(off) as *const __m256i);
5010                let p16 = _mm256_maddubs_epi16(aw, _mm256_sign_epi8(x, w));
5011                _mm256_cvtepi32_ps(_mm256_madd_epi16(p16, ones))
5012            };
5013            f0 = _mm256_fmadd_ps(dot(xs[0]), sv, f0);
5014            f1 = _mm256_fmadd_ps(dot(xs[1]), sv, f1);
5015            f2 = _mm256_fmadd_ps(dot(xs[2]), sv, f2);
5016            f3 = _mm256_fmadd_ps(dot(xs[3]), sv, f3);
5017        }
5018        [
5019            hsum256_ps(f0),
5020            hsum256_ps(f1),
5021            hsum256_ps(f2),
5022            hsum256_ps(f3),
5023        ]
5024    }
5025}
5026
5027/// Horizontal sum of eight f32 lanes — the one cross-lane reduction the
5028/// blocked kernels pay, once per row instead of once per group.
5029#[cfg(target_arch = "x86_64")]
5030#[target_feature(enable = "avx2")]
5031#[inline]
5032unsafe fn hsum256_ps(v: core::arch::x86_64::__m256) -> f32 {
5033    // SAFETY: pure register arithmetic on the caller's vector.
5034    unsafe {
5035        use core::arch::x86_64::*;
5036        let hi = _mm256_extractf128_ps::<1>(v);
5037        let s = _mm_add_ps(_mm256_castps256_ps128(v), hi);
5038        let s = _mm_add_ps(s, _mm_movehl_ps(s, s));
5039        let s = _mm_add_ss(s, _mm_shuffle_ps::<0x55>(s, s));
5040        _mm_cvtss_f32(s)
5041    }
5042}
5043
5044/// VNNI twin of `dot_q4t_row_1x4_avx2` (see `dpbusd_hsum`).
5045#[cfg(target_arch = "x86_64")]
5046#[target_feature(enable = "avx2,fma,avx512f,avx512bw,avx512vl,avx512vnni")]
5047unsafe fn dot_q4t_row_1x4_vnni(bytes: &[u8], r: usize, gpr: usize, xs: [&[i8]; 4]) -> [f32; 4] {
5048    // SAFETY: callers uphold the 18B-tile and xq-length contracts.
5049    unsafe {
5050        use core::arch::x86_64::*;
5051        let lomask = _mm_set1_epi8(0x0F);
5052        let eight = _mm256_set1_epi8(8);
5053        // Same shape as the AVX2 twin: accumulate in f32 vectors and pay
5054        // one cross-lane reduction per row, not per (group, activation).
5055        let mut f0 = _mm256_setzero_ps();
5056        let mut f1 = _mm256_setzero_ps();
5057        let mut f2 = _mm256_setzero_ps();
5058        let mut f3 = _mm256_setzero_ps();
5059        for gi in 0..gpr {
5060            let t = bytes.as_ptr().add((r * gpr + gi) * Q4_TILE);
5061            let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
5062            let sv = _mm256_set1_ps(s);
5063            let bb = _mm_loadu_si128(t.add(2) as *const __m128i);
5064            let lo = _mm_and_si128(bb, lomask);
5065            let hi = _mm_and_si128(_mm_srli_epi16::<4>(bb), lomask);
5066            let w = _mm256_sub_epi8(
5067                _mm256_set_m128i(_mm_unpackhi_epi8(lo, hi), _mm_unpacklo_epi8(lo, hi)),
5068                eight,
5069            );
5070            let aw = _mm256_abs_epi8(w);
5071            let off = gi * GROUP_SIZE;
5072            let dot = |xq: &[i8]| {
5073                let x = _mm256_loadu_si256(xq.as_ptr().add(off) as *const __m256i);
5074                _mm256_cvtepi32_ps(_mm256_dpbusd_epi32(
5075                    _mm256_setzero_si256(),
5076                    aw,
5077                    _mm256_sign_epi8(x, w),
5078                ))
5079            };
5080            f0 = _mm256_fmadd_ps(dot(xs[0]), sv, f0);
5081            f1 = _mm256_fmadd_ps(dot(xs[1]), sv, f1);
5082            f2 = _mm256_fmadd_ps(dot(xs[2]), sv, f2);
5083            f3 = _mm256_fmadd_ps(dot(xs[3]), sv, f3);
5084        }
5085        let acc = [
5086            hsum256_ps(f0),
5087            hsum256_ps(f1),
5088            hsum256_ps(f2),
5089            hsum256_ps(f3),
5090        ];
5091        acc
5092    }
5093}
5094
5095/// ARM twin of `dot_q4t_row_1x4_avx2`: one nibble unpack per group
5096/// serves FOUR activation streams. Per stream the group order and f32
5097/// accumulation match `dot_q4t_row_sdot` exactly — batch == matvec
5098/// bit-for-bit.
5099#[cfg(target_arch = "aarch64")]
5100#[target_feature(enable = "neon,dotprod")]
5101unsafe fn dot_q4t_row_1x4_sdot(bytes: &[u8], r: usize, gpr: usize, xs: [&[i8]; 4]) -> [f32; 4] {
5102    // SAFETY: callers uphold the 18B-tile and xq-length contracts.
5103    unsafe {
5104        use core::arch::aarch64::*;
5105        use core::arch::asm;
5106        let lomask = vdupq_n_u8(0x0F);
5107        let eight = vdupq_n_s8(8);
5108        let mut acc = [0f32; 4];
5109        for gi in 0..gpr {
5110            let t = bytes.as_ptr().add((r * gpr + gi) * Q4_TILE);
5111            let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
5112            let b = vld1q_u8(t.add(2));
5113            let lo = vandq_u8(b, lomask);
5114            let hi = vshrq_n_u8::<4>(b);
5115            let e0 = vsubq_s8(vreinterpretq_s8_u8(vzip1q_u8(lo, hi)), eight);
5116            let e1 = vsubq_s8(vreinterpretq_s8_u8(vzip2q_u8(lo, hi)), eight);
5117            for (k, xq) in xs.iter().enumerate() {
5118                let x0 = vld1q_s8(xq.as_ptr().add(gi * GROUP_SIZE));
5119                let x1 = vld1q_s8(xq.as_ptr().add(gi * GROUP_SIZE + 16));
5120                let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
5121                asm!(
5122                    "sdot {a0:v}.4s, {e0:v}.16b, {x0:v}.16b",
5123                    "sdot {a1:v}.4s, {e1:v}.16b, {x1:v}.16b",
5124                    a0 = inout(vreg) a0, a1 = inout(vreg) a1,
5125                    e0 = in(vreg) e0, x0 = in(vreg) x0, e1 = in(vreg) e1, x1 = in(vreg) x1,
5126                    options(pure, nomem, nostack),
5127                );
5128                acc[k] += vaddvq_s32(vaddq_s32(a0, a1)) as f32 * s;
5129            }
5130        }
5131        acc
5132    }
5133}
5134
5135/// Exact-term correction for A8W8 outliers on a tiled row.
5136#[inline]
5137fn q4t_outlier(bytes: &[u8], r: usize, gpr: usize, j: usize) -> (f32, f32) {
5138    let gi = j / GROUP_SIZE;
5139    let k = j % GROUP_SIZE;
5140    let tile = &bytes[(r * gpr + gi) * Q4_TILE..(r * gpr + gi + 1) * Q4_TILE];
5141    let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
5142    let byte = tile[2 + k / 2];
5143    let nib = if k & 1 == 0 { byte & 0x0F } else { byte >> 4 };
5144    ((nib as i32 - 8) as f32, s)
5145}
5146
5147/// Exact scalar q4_tiled row (CMF_SDOT=0 contract) — same pairwise
5148/// accumulation shape as `q4_range_f32`.
5149#[inline]
5150fn q4t_row_exact(bytes: &[u8], r: usize, gpr: usize, x: &[f32]) -> f32 {
5151    let mut acc = 0f32;
5152    for gi in 0..gpr {
5153        let tile = &bytes[(r * gpr + gi) * Q4_TILE..(r * gpr + gi + 1) * Q4_TILE];
5154        let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
5155        let xg = &x[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
5156        let mut ga = 0f32;
5157        for (k, &b) in tile[2..].iter().enumerate() {
5158            ga += ((b & 0x0F) as f32 - 8.0) * xg[k * 2]
5159                + (((b >> 4) & 0x0F) as f32 - 8.0) * xg[k * 2 + 1];
5160        }
5161        acc += ga * s;
5162    }
5163    acc
5164}
5165
5166/// Split view of a `q4tp` payload. The three planes are resolved once per
5167/// matvec instead of per row — `q4tp_sections` is cheap, but doing it inside
5168/// the row loop would put a division on the hot path for nothing.
5169struct Q4tpView<'a> {
5170    nib: &'a [u8],
5171    params: &'a [u8],
5172    codes: &'a [u8],
5173    stride: usize,
5174    /// q2tp reads the ladder with rung 0 = exact zero.
5175    zero_rung: bool,
5176}
5177
5178impl<'a> Q4tpView<'a> {
5179    fn new(bytes: &'a [u8], rows: usize, cols: usize) -> Self {
5180        let (params_off, codes_off, stride) = q4tp_sections(rows, cols);
5181        Self {
5182            nib: &bytes[..params_off],
5183            params: &bytes[params_off..codes_off],
5184            codes: &bytes[codes_off..],
5185            stride,
5186            zero_rung: false,
5187        }
5188    }
5189
5190    /// The q2tp view: identical params/codes planes, 8 B weight chunks.
5191    fn new_q2(bytes: &'a [u8], rows: usize, cols: usize) -> Self {
5192        let (params_off, codes_off, stride) = q2tp_sections(rows, cols);
5193        Self {
5194            nib: &bytes[..params_off],
5195            params: &bytes[params_off..codes_off],
5196            codes: &bytes[codes_off..],
5197            stride,
5198            zero_rung: true,
5199        }
5200    }
5201
5202    /// Expand row `r`'s per-tile scales into `out` (length `gpr`).
5203    ///
5204    /// Doing this once per row — rather than decoding a 5-bit code inside the
5205    /// tile loop — is what makes the format free at runtime. Random access to
5206    /// a packed 5-bit field costs a division, two bounds checks and a branch;
5207    /// the tile's actual work is two `sdot`s, so per-tile decoding dominated
5208    /// the kernel and cost 5x (measured: 1.4 vs 6.9 tok/s on Nanbeige-3B).
5209    /// Walking the plane sequentially with a bit accumulator is ~3 ops.
5210    /// Eight 5-bit codes are exactly five bytes, so a whole group of
5211    /// eight decodes from one little-endian word at fixed shifts. The
5212    /// bit-accumulator this replaces carried a data-dependent `while
5213    /// have < 5` refill whose branch sat in the innermost loop of every
5214    /// q4tp row; a decode profile put this function above the dot
5215    /// products it feeds. Same bitstream, same codes — just no branch
5216    /// and eight independent extractions.
5217    #[inline]
5218    fn scales_into(&self, r: usize, gpr: usize, out: &mut [f32]) {
5219        let tab = if self.zero_rung {
5220            q2tp_ladder(self.params, r)
5221        } else {
5222            q4tp_ladder(self.params, r)
5223        };
5224        let codes = &self.codes[r * self.stride..(r + 1) * self.stride];
5225        let out = &mut out[..gpr];
5226        let mut chunks = out.chunks_exact_mut(8);
5227        let mut ci = 0usize;
5228        for c in &mut chunks {
5229            let w = u64::from(codes[ci])
5230                | u64::from(codes[ci + 1]) << 8
5231                | u64::from(codes[ci + 2]) << 16
5232                | u64::from(codes[ci + 3]) << 24
5233                | u64::from(codes[ci + 4]) << 32;
5234            for (k, o) in c.iter_mut().enumerate() {
5235                *o = tab[((w >> (5 * k)) & 31) as usize];
5236            }
5237            ci += 5;
5238        }
5239        // Fewer than eight codes left: the shared total accessor, which
5240        // tolerates a 5-bit field whose spill byte is past the stride.
5241        let tail = &codes[ci..];
5242        for (k, o) in chunks.into_remainder().iter_mut().enumerate() {
5243            *o = tab[q4tp_code(tail, k)];
5244        }
5245    }
5246}
5247
5248#[inline]
5249fn dot_q4tp_row_i8(nib: &[u8], r: usize, gpr: usize, xq: &[i8], scales: &[f32]) -> f32 {
5250    #[cfg(target_arch = "aarch64")]
5251    unsafe {
5252        return dot_q4tp_row_sdot(nib, r, gpr, xq, scales);
5253    }
5254    #[cfg(target_arch = "x86_64")]
5255    unsafe {
5256        if vnni_tiles_enabled() {
5257            return dot_q4tp_row_vnni(nib, r, gpr, xq, scales);
5258        }
5259        return dot_q4tp_row_avx2(nib, r, gpr, xq, scales);
5260    }
5261    #[allow(unreachable_code)]
5262    {
5263        let mut acc = 0f32;
5264        for gi in 0..gpr {
5265            let tile = &nib[(r * gpr + gi) * Q4TP_NIB..(r * gpr + gi + 1) * Q4TP_NIB];
5266            let s = scales[gi];
5267            let mut d = 0i32;
5268            for (k, &b) in tile.iter().enumerate() {
5269                d += ((b & 0x0F) as i32 - 8) * xq[gi * GROUP_SIZE + k * 2] as i32
5270                    + (((b >> 4) & 0x0F) as i32 - 8) * xq[gi * GROUP_SIZE + k * 2 + 1] as i32;
5271            }
5272            acc += d as f32 * s;
5273        }
5274        acc
5275    }
5276}
5277
5278/// q4tp twin of `dot_q4t_row_sdot`: identical nibble math, but the tile
5279/// stride is 16 B (no inline scale) and the scale is a ladder lookup.
5280#[cfg(target_arch = "aarch64")]
5281#[target_feature(enable = "neon,dotprod")]
5282unsafe fn dot_q4tp_row_sdot(nib: &[u8], r: usize, gpr: usize, xq: &[i8], scales: &[f32]) -> f32 {
5283    // SAFETY: callers uphold slice-length contracts (16B tile per group,
5284    // xq.len() == gpr·GROUP_SIZE, codes covering gpr 5-bit fields).
5285    unsafe {
5286        use core::arch::aarch64::*;
5287        use core::arch::asm;
5288        let lomask = vdupq_n_u8(0x0F);
5289        let eight = vdupq_n_s8(8);
5290        let mut acc = 0f32;
5291        for gi in 0..gpr {
5292            let t = nib.as_ptr().add((r * gpr + gi) * Q4TP_NIB);
5293            let s = *scales.get_unchecked(gi);
5294            let b = vld1q_u8(t);
5295            let lo = vandq_u8(b, lomask);
5296            let hi = vshrq_n_u8::<4>(b);
5297            let e0 = vsubq_s8(vreinterpretq_s8_u8(vzip1q_u8(lo, hi)), eight);
5298            let e1 = vsubq_s8(vreinterpretq_s8_u8(vzip2q_u8(lo, hi)), eight);
5299            let x0 = vld1q_s8(xq.as_ptr().add(gi * GROUP_SIZE));
5300            let x1 = vld1q_s8(xq.as_ptr().add(gi * GROUP_SIZE + 16));
5301            let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
5302            asm!(
5303                "sdot {a0:v}.4s, {e0:v}.16b, {x0:v}.16b",
5304                "sdot {a1:v}.4s, {e1:v}.16b, {x1:v}.16b",
5305                a0 = inout(vreg) a0, a1 = inout(vreg) a1,
5306                e0 = in(vreg) e0, x0 = in(vreg) x0, e1 = in(vreg) e1, x1 = in(vreg) x1,
5307                options(pure, nomem, nostack),
5308            );
5309            acc += vaddvq_s32(vaddq_s32(a0, a1)) as f32 * s;
5310        }
5311        acc
5312    }
5313}
5314
5315#[cfg(target_arch = "x86_64")]
5316#[target_feature(enable = "avx2")]
5317unsafe fn dot_q4tp_row_avx2(nib: &[u8], r: usize, gpr: usize, xq: &[i8], scales: &[f32]) -> f32 {
5318    // SAFETY: see dot_q4tp_row_sdot.
5319    unsafe {
5320        use core::arch::x86_64::*;
5321        let lomask = _mm_set1_epi8(0x0F);
5322        let eight = _mm256_set1_epi8(8);
5323        let ones = _mm256_set1_epi16(1);
5324        let mut acc = 0f32;
5325        for gi in 0..gpr {
5326            let t = nib.as_ptr().add((r * gpr + gi) * Q4TP_NIB);
5327            let s = *scales.get_unchecked(gi);
5328            let b = _mm_loadu_si128(t as *const __m128i);
5329            let lo = _mm_and_si128(b, lomask);
5330            let hi = _mm_and_si128(_mm_srli_epi16::<4>(b), lomask);
5331            let w = _mm256_sub_epi8(
5332                _mm256_set_m128i(_mm_unpackhi_epi8(lo, hi), _mm_unpacklo_epi8(lo, hi)),
5333                eight,
5334            );
5335            let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
5336            let p16 = _mm256_maddubs_epi16(_mm256_abs_epi8(w), _mm256_sign_epi8(x, w));
5337            let d = _mm256_madd_epi16(p16, ones);
5338            let hi128 = _mm256_extracti128_si256::<1>(d);
5339            let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
5340            let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
5341            let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
5342            acc += _mm_cvtsi128_si32(s32) as f32 * s;
5343        }
5344        acc
5345    }
5346}
5347
5348/// VNNI twin of `dot_q4tp_row_avx2` (see `dot_q4t_row_vnni` for why the
5349/// 256-bit VL encoding is the one to use here).
5350#[cfg(target_arch = "x86_64")]
5351#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
5352unsafe fn dot_q4tp_row_vnni(nib: &[u8], r: usize, gpr: usize, xq: &[i8], scales: &[f32]) -> f32 {
5353    // SAFETY: see dot_q4tp_row_sdot.
5354    unsafe {
5355        use core::arch::x86_64::*;
5356        let lomask = _mm_set1_epi8(0x0F);
5357        let eight = _mm256_set1_epi8(8);
5358        let mut acc = 0f32;
5359        for gi in 0..gpr {
5360            let t = nib.as_ptr().add((r * gpr + gi) * Q4TP_NIB);
5361            let s = *scales.get_unchecked(gi);
5362            let b = _mm_loadu_si128(t as *const __m128i);
5363            let lo = _mm_and_si128(b, lomask);
5364            let hi = _mm_and_si128(_mm_srli_epi16::<4>(b), lomask);
5365            let w = _mm256_sub_epi8(
5366                _mm256_set_m128i(_mm_unpackhi_epi8(lo, hi), _mm_unpacklo_epi8(lo, hi)),
5367                eight,
5368            );
5369            let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
5370            acc += dpbusd_hsum(_mm256_abs_epi8(w), _mm256_sign_epi8(x, w)) as f32 * s;
5371        }
5372        acc
5373    }
5374}
5375
5376/// Exact scalar q4tp row — the `CMF_SDOT=0` contract, same pairwise
5377/// accumulation shape as `q4t_row_exact`.
5378#[inline]
5379fn q4tp_row_exact(nib: &[u8], r: usize, gpr: usize, x: &[f32], scales: &[f32]) -> f32 {
5380    #[cfg(target_arch = "x86_64")]
5381    if avx2_enabled() {
5382        // Keep the scalar pair/group reduction order, not a vector sum.
5383        return unsafe { q4tp_row_float_avx2(nib, r, gpr, x, scales) };
5384    }
5385    q4tp_row_float_scalar(nib, r, gpr, x, scales)
5386}
5387
5388#[inline]
5389fn q4tp_row_float_scalar(nib: &[u8], r: usize, gpr: usize, x: &[f32], scales: &[f32]) -> f32 {
5390    let mut acc = 0f32;
5391    for gi in 0..gpr {
5392        let tile = &nib[(r * gpr + gi) * Q4TP_NIB..(r * gpr + gi + 1) * Q4TP_NIB];
5393        let s = scales[gi];
5394        let xg = &x[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
5395        let mut ga = 0f32;
5396        for (k, &b) in tile.iter().enumerate() {
5397            ga += ((b & 0x0F) as f32 - 8.0) * xg[k * 2]
5398                + (((b >> 4) & 0x0F) as f32 - 8.0) * xg[k * 2 + 1];
5399        }
5400        acc += ga * s;
5401    }
5402    acc
5403}
5404
5405/// Vectorize unpack, conversion and multiplication, but preserve every
5406/// pair addition and the scalar accumulation order. No activation rounding
5407/// or FMA: bit-identical to the float scalar row, including its group scale.
5408#[cfg(target_arch = "x86_64")]
5409#[target_feature(enable = "avx2")]
5410unsafe fn q4tp_row_float_avx2(nib: &[u8], r: usize, gpr: usize, x: &[f32], scales: &[f32]) -> f32 {
5411    // SAFETY: caller checks AVX2 and provides the same complete 32-element
5412    // groups as the scalar row. Loads/stores are explicitly unaligned.
5413    unsafe {
5414        use core::arch::x86_64::*;
5415        let mask = _mm_set1_epi8(15);
5416        let eight = _mm_set1_epi8(8);
5417        let order = _mm256_setr_epi32(0, 1, 4, 5, 2, 3, 6, 7);
5418        let mut acc = 0.0f32;
5419        for gi in 0..gpr {
5420            let packed = _mm_loadu_si128(nib.as_ptr().add((r * gpr + gi) * Q4TP_NIB).cast());
5421            let lo = _mm_and_si128(packed, mask);
5422            let hi = _mm_and_si128(_mm_srli_epi16::<4>(packed), mask);
5423            let w0 = _mm_sub_epi8(_mm_unpacklo_epi8(lo, hi), eight);
5424            let w1 = _mm_sub_epi8(_mm_unpackhi_epi8(lo, hi), eight);
5425            let xp = x.as_ptr().add(gi * GROUP_SIZE);
5426            let a = _mm256_mul_ps(
5427                _mm256_cvtepi32_ps(_mm256_cvtepi8_epi32(w0)),
5428                _mm256_loadu_ps(xp),
5429            );
5430            let b = _mm256_mul_ps(
5431                _mm256_cvtepi32_ps(_mm256_cvtepi8_epi32(_mm_srli_si128::<8>(w0))),
5432                _mm256_loadu_ps(xp.add(8)),
5433            );
5434            let c = _mm256_mul_ps(
5435                _mm256_cvtepi32_ps(_mm256_cvtepi8_epi32(w1)),
5436                _mm256_loadu_ps(xp.add(16)),
5437            );
5438            let d = _mm256_mul_ps(
5439                _mm256_cvtepi32_ps(_mm256_cvtepi8_epi32(_mm_srli_si128::<8>(w1))),
5440                _mm256_loadu_ps(xp.add(24)),
5441            );
5442            let mut pairs = [0.0f32; 16];
5443            _mm256_storeu_ps(
5444                pairs.as_mut_ptr(),
5445                _mm256_permutevar8x32_ps(_mm256_hadd_ps(a, b), order),
5446            );
5447            _mm256_storeu_ps(
5448                pairs.as_mut_ptr().add(8),
5449                _mm256_permutevar8x32_ps(_mm256_hadd_ps(c, d), order),
5450            );
5451            let mut ga = 0.0f32;
5452            for v in pairs {
5453                ga += v;
5454            }
5455            acc += ga * scales[gi];
5456        }
5457        acc
5458    }
5459}
5460
5461/// Single weight of a q4tp tensor — the a8w8 outlier path, which restores
5462/// activation outliers at full precision after the int8 pass.
5463#[inline]
5464fn q4tp_outlier(nib: &[u8], r: usize, gpr: usize, j: usize, scales: &[f32]) -> (f32, f32) {
5465    let (gi, k) = (j / GROUP_SIZE, j % GROUP_SIZE);
5466    let byte = nib[(r * gpr + gi) * Q4TP_NIB + k / 2];
5467    let n = if k & 1 == 0 { byte & 0x0F } else { byte >> 4 };
5468    ((n as i32 - 8) as f32, scales[gi])
5469}
5470
5471/// Fused q4tp matvec (dispatch mirrors `q4t_matvec`).
5472fn q4tp_matvec(
5473    bytes: &[u8],
5474    x: &[f32],
5475    rows: usize,
5476    cols: usize,
5477    out: &mut [f32],
5478    pool: Option<&Pool>,
5479) {
5480    debug_assert_eq!(out.len(), rows);
5481    let gpr = cols / GROUP_SIZE;
5482    let v = Q4tpView::new(bytes, rows, cols);
5483    let out_addr = SendMut(out.as_mut_ptr());
5484    if a8w8_enabled() {
5485        let act = split_act(x);
5486        let run = |start: usize, end: usize| {
5487            // One scratch row of scales per worker — borrowed, not minted.
5488            with_krow(gpr, |sc| {
5489                for r in start..end {
5490                    v.scales_into(r, gpr, sc);
5491                    let mut acc = dot_q4tp_row_i8(v.nib, r, gpr, &act.xq, sc) * act.sx;
5492                    for &(j, xv) in &act.outliers {
5493                        let (w, s) = q4tp_outlier(v.nib, r, gpr, j, sc);
5494                        acc += w * s * xv;
5495                    }
5496                    // SAFETY: disjoint row ranges per worker.
5497                    unsafe { *out_addr.at(r) = acc };
5498                }
5499            })
5500        };
5501        dispatch_rows(pool, rows, &run);
5502        return;
5503    }
5504    let run = |start: usize, end: usize| {
5505        with_krow(gpr, |sc| {
5506            for r in start..end {
5507                v.scales_into(r, gpr, sc);
5508                // SAFETY: disjoint row ranges per worker.
5509                unsafe { *out_addr.at(r) = q4tp_row_exact(v.nib, r, gpr, x, sc) };
5510            }
5511        })
5512    };
5513    dispatch_rows(pool, rows, &run);
5514}
5515
5516/// Fused two-input q4tp matvec — the SwiGLU gate/up pair. Weights and the
5517/// row ladder are read once and spent on both activation streams.
5518#[allow(clippy::too_many_arguments)]
5519fn q4tp_matvec2(
5520    bytes: &[u8],
5521    x1: &[f32],
5522    x2: &[f32],
5523    rows: usize,
5524    cols: usize,
5525    o1: &mut [f32],
5526    o2: &mut [f32],
5527    pool: Option<&Pool>,
5528) {
5529    let gpr = cols / GROUP_SIZE;
5530    let v = Q4tpView::new(bytes, rows, cols);
5531    let (p1, p2) = (SendMut(o1.as_mut_ptr()), SendMut(o2.as_mut_ptr()));
5532    let run = |start: usize, end: usize| {
5533        let mut sc = vec![0f32; gpr];
5534        for r in start..end {
5535            v.scales_into(r, gpr, &mut sc);
5536            // SAFETY: disjoint row ranges per worker.
5537            unsafe {
5538                *p1.at(r) = q4tp_row_exact(v.nib, r, gpr, x1, &sc);
5539                *p2.at(r) = q4tp_row_exact(v.nib, r, gpr, x2, &sc);
5540            }
5541        }
5542    };
5543    dispatch_rows(pool, rows, &run);
5544}
5545
5546/// One q2tp outlier weight at column `j` of row `r`: the 2-bit code and
5547/// its group scale, mirrored on `q4tp_outlier`.
5548#[inline]
5549fn q2tp_outlier(chunks: &[u8], r: usize, gpr: usize, j: usize, scales: &[f32]) -> (f32, f32) {
5550    let (gi, k) = (j / GROUP_SIZE, j % GROUP_SIZE);
5551    let byte = chunks[(r * gpr + gi) * Q2TP_CHUNK + k / 4];
5552    let c = (byte >> (2 * (k % 4))) & 3;
5553    (c as f32 - 1.5, scales[gi])
5554}
5555
5556#[cfg(target_arch = "x86_64")]
5557const Q2TP_DECODE_U32: [u32; 256] = {
5558    let mut tab = [0u32; 256];
5559    let mut b = 0usize;
5560    while b < 256 {
5561        tab[b] = ((b as u32) & 3)
5562            | ((((b as u32) >> 2) & 3) << 8)
5563            | ((((b as u32) >> 4) & 3) << 16)
5564            | ((((b as u32) >> 6) & 3) << 24);
5565        b += 1;
5566    }
5567    tab
5568};
5569
5570/// Eight packed q2tp bytes against 32 signed activation bytes. `maddubs`
5571/// exactly computes unsigned 2-bit code × signed i8; its pair sums cannot
5572/// saturate (2 × 3 × 127 < i16::MAX), and the second madd widens to i32.
5573#[cfg(target_arch = "x86_64")]
5574#[target_feature(enable = "avx2")]
5575unsafe fn q2tp_code_dot_avx2(ch: &[u8], x: &[i8]) -> i32 {
5576    use core::arch::x86_64::*;
5577    debug_assert!(ch.len() >= Q2TP_CHUNK && x.len() >= GROUP_SIZE);
5578    let codes = _mm256_setr_epi32(
5579        Q2TP_DECODE_U32[ch[0] as usize] as i32,
5580        Q2TP_DECODE_U32[ch[1] as usize] as i32,
5581        Q2TP_DECODE_U32[ch[2] as usize] as i32,
5582        Q2TP_DECODE_U32[ch[3] as usize] as i32,
5583        Q2TP_DECODE_U32[ch[4] as usize] as i32,
5584        Q2TP_DECODE_U32[ch[5] as usize] as i32,
5585        Q2TP_DECODE_U32[ch[6] as usize] as i32,
5586        Q2TP_DECODE_U32[ch[7] as usize] as i32,
5587    );
5588    let xv = unsafe { _mm256_loadu_si256(x.as_ptr().cast()) };
5589    let pair = _mm256_maddubs_epi16(codes, xv);
5590    let quad = _mm256_madd_epi16(pair, _mm256_set1_epi16(1));
5591    let sum128 = _mm_add_epi32(
5592        _mm256_castsi256_si128(quad),
5593        _mm256_extracti128_si256(quad, 1),
5594    );
5595    let sum64 = _mm_hadd_epi32(sum128, sum128);
5596    _mm_cvtsi128_si32(_mm_hadd_epi32(sum64, sum64))
5597}
5598
5599/// Integer dot of one q2tp row against pre-quantized activations:
5600/// Σ_g s_g · (Σ c·xq − 1.5·Σ xq). The half-integer grid (c − 1.5)
5601/// becomes exact integer math through the group sums — the same trick
5602/// every a8w8 kernel in this file rides. The codes decode into a
5603/// 32-byte scratch in natural order and the dot itself is the shared
5604/// SDOT primitive; elsewhere a scalar integer loop.
5605#[inline]
5606fn dot_q2tp_row_i8(
5607    chunks: &[u8],
5608    r: usize,
5609    gpr: usize,
5610    xq: &[i8],
5611    gsum: &[i32],
5612    scales: &[f32],
5613) -> f32 {
5614    let mut acc = 0f32;
5615    let base = r * gpr * Q2TP_CHUNK;
5616    #[cfg(not(any(target_arch = "aarch64", target_arch = "x86_64")))]
5617    let mut codes = [0i8; GROUP_SIZE];
5618    #[cfg(target_arch = "x86_64")]
5619    let avx2 = std::arch::is_x86_feature_detected!("avx2");
5620    for gi in 0..gpr {
5621        let ch = &chunks[base + gi * Q2TP_CHUNK..base + (gi + 1) * Q2TP_CHUNK];
5622        let xg = &xq[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
5623        #[cfg(target_arch = "aarch64")]
5624        // NEON: the byte's four 2-bit fields land in four lane vectors
5625        // (shift+mask), vld4 de-interleaves xq to match (xj[k] =
5626        // xq[4k+j]), widening MACs accumulate exactly in i32. A scalar
5627        // decode here cost as much as the dot it fed — the profile put
5628        // it at the top of the whole W2 decode.
5629        let dot = unsafe {
5630            use core::arch::aarch64::*;
5631            let b = vld1_u8(ch.as_ptr());
5632            let three = vdup_n_u8(3);
5633            let c0 = vreinterpret_s8_u8(vand_u8(b, three));
5634            let c1 = vreinterpret_s8_u8(vand_u8(vshr_n_u8(b, 2), three));
5635            let c2 = vreinterpret_s8_u8(vand_u8(vshr_n_u8(b, 4), three));
5636            let c3 = vreinterpret_s8_u8(vand_u8(vshr_n_u8(b, 6), three));
5637            let x4 = vld4_s8(xg.as_ptr());
5638            let mut acc4 = vdupq_n_s32(0);
5639            acc4 = vpadalq_s16(acc4, vmull_s8(c0, x4.0));
5640            acc4 = vpadalq_s16(acc4, vmull_s8(c1, x4.1));
5641            acc4 = vpadalq_s16(acc4, vmull_s8(c2, x4.2));
5642            acc4 = vpadalq_s16(acc4, vmull_s8(c3, x4.3));
5643            vaddvq_s32(acc4)
5644        };
5645        #[cfg(target_arch = "x86_64")]
5646        let dot: i32 = if avx2 {
5647            // SAFETY: the runtime feature check gates the target-feature body;
5648            // the group slices above are exactly 8 and 32 bytes long.
5649            unsafe { q2tp_code_dot_avx2(ch, xg) }
5650        } else {
5651            ch.iter()
5652                .enumerate()
5653                .map(|(k, &b)| {
5654                    ((b & 3) as i32) * xg[k * 4] as i32
5655                        + (((b >> 2) & 3) as i32) * xg[k * 4 + 1] as i32
5656                        + (((b >> 4) & 3) as i32) * xg[k * 4 + 2] as i32
5657                        + (((b >> 6) & 3) as i32) * xg[k * 4 + 3] as i32
5658                })
5659                .sum()
5660        };
5661        #[cfg(not(any(target_arch = "aarch64", target_arch = "x86_64")))]
5662        let dot: i32 = {
5663            for (k, &b) in ch.iter().enumerate() {
5664                codes[k * 4] = (b & 3) as i8;
5665                codes[k * 4 + 1] = ((b >> 2) & 3) as i8;
5666                codes[k * 4 + 2] = ((b >> 4) & 3) as i8;
5667                codes[k * 4 + 3] = ((b >> 6) & 3) as i8;
5668            }
5669            codes
5670                .iter()
5671                .zip(xg)
5672                .map(|(&c, &x)| c as i32 * x as i32)
5673                .sum()
5674        };
5675        acc += scales[gi] * (dot as f32 - 1.5 * gsum[gi] as f32);
5676    }
5677    acc
5678}
5679
5680/// Exact f32 dot of one q2tp row: 2-bit fields LSB-first, (c − 1.5)·s.
5681/// Scalar on purpose — the 2-bit class targets the GPU graph; the CPU
5682/// path exists for parity gates and small-machine fallback.
5683fn q2tp_row_exact(chunks: &[u8], r: usize, gpr: usize, x: &[f32], scales: &[f32]) -> f32 {
5684    q2tp_row_exact_center(chunks, r, gpr, x, scales, 1.5)
5685}
5686
5687/// Fused Prism affine row: the derived correction is applied inside the
5688/// decoded code, avoiding a second accumulated dot and avoiding cancellation
5689/// between `B=(c-1.5)s` and `+.5s` for long 5120/17408 rows.
5690#[inline]
5691fn q2tp_affine_row_exact(chunks: &[u8], r: usize, gpr: usize, x: &[f32], scales: &[f32]) -> f32 {
5692    q2tp_row_exact_center(chunks, r, gpr, x, scales, 1.0)
5693}
5694
5695#[inline]
5696fn q2tp_row_exact_center(
5697    chunks: &[u8],
5698    r: usize,
5699    gpr: usize,
5700    x: &[f32],
5701    scales: &[f32],
5702    center: f32,
5703) -> f32 {
5704    let mut acc = 0f32;
5705    for gi in 0..gpr {
5706        let ch = &chunks[(r * gpr + gi) * Q2TP_CHUNK..(r * gpr + gi + 1) * Q2TP_CHUNK];
5707        let s = scales[gi];
5708        let xb = &x[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
5709        let mut g = 0f32;
5710        for (k, &b) in ch.iter().enumerate() {
5711            g += ((b & 3) as f32 - center) * xb[k * 4]
5712                + (((b >> 2) & 3) as f32 - center) * xb[k * 4 + 1]
5713                + (((b >> 4) & 3) as f32 - center) * xb[k * 4 + 2]
5714                + (((b >> 6) & 3) as f32 - center) * xb[k * 4 + 3];
5715        }
5716        acc += s * g;
5717    }
5718    acc
5719}
5720
5721fn q2tp_matvec(
5722    bytes: &[u8],
5723    x: &[f32],
5724    rows: usize,
5725    cols: usize,
5726    out: &mut [f32],
5727    pool: Option<&Pool>,
5728) {
5729    q2tp_matvec_mode(bytes, x, rows, cols, out, pool, false);
5730}
5731
5732fn q2tp_affine_matvec(
5733    bytes: &[u8],
5734    x: &[f32],
5735    rows: usize,
5736    cols: usize,
5737    out: &mut [f32],
5738    pool: Option<&Pool>,
5739) {
5740    q2tp_matvec_mode(bytes, x, rows, cols, out, pool, true);
5741}
5742
5743fn q2tp_matvec_mode(
5744    bytes: &[u8],
5745    x: &[f32],
5746    rows: usize,
5747    cols: usize,
5748    out: &mut [f32],
5749    pool: Option<&Pool>,
5750    affine: bool,
5751) {
5752    debug_assert_eq!(out.len(), rows);
5753    let gpr = cols / GROUP_SIZE;
5754    let v = Q4tpView::new_q2(bytes, rows, cols);
5755    let out_addr = SendMut(out.as_mut_ptr());
5756    // a8w8 fast path (CMF_SDOT=0 keeps the exact scalar walk): integer
5757    // code dots + group sums, exact outlier correction — the same
5758    // contract as every sibling kernel; measured 2-bit rows were the
5759    // only scalar holdout in the family.
5760    if !affine && a8w8_enabled() {
5761        let act = split_act(x);
5762        let gsum = q1_group_sums(&act.xq, gpr);
5763        let (act, gsum) = (&act, &gsum);
5764        let run = move |start: usize, end: usize| {
5765            with_krow(gpr, |sc| {
5766                for r in start..end {
5767                    v.scales_into(r, gpr, sc);
5768                    let mut acc = dot_q2tp_row_i8(v.nib, r, gpr, &act.xq, gsum, sc) * act.sx;
5769                    for &(j, xv) in &act.outliers {
5770                        let (w, s) = q2tp_outlier(v.nib, r, gpr, j, sc);
5771                        acc += w * s * xv;
5772                    }
5773                    // SAFETY: disjoint row ranges per worker.
5774                    unsafe { *out_addr.at(r) = acc };
5775                }
5776            })
5777        };
5778        dispatch_rows(pool, rows, &run);
5779        return;
5780    }
5781    let run = |start: usize, end: usize| {
5782        with_krow(gpr, |sc| {
5783            for r in start..end {
5784                v.scales_into(r, gpr, sc);
5785                // SAFETY: disjoint row ranges per worker.
5786                unsafe {
5787                    *out_addr.at(r) = if affine {
5788                        q2tp_affine_row_exact(v.nib, r, gpr, x, sc)
5789                    } else {
5790                        q2tp_row_exact(v.nib, r, gpr, x, sc)
5791                    }
5792                };
5793            }
5794        })
5795    };
5796    dispatch_rows(pool, rows, &run);
5797}
5798
5799/// Fused two-input q2tp matvec — the SwiGLU gate/up pair.
5800#[allow(clippy::too_many_arguments)]
5801fn q2tp_matvec2(
5802    bytes: &[u8],
5803    x1: &[f32],
5804    x2: &[f32],
5805    rows: usize,
5806    cols: usize,
5807    o1: &mut [f32],
5808    o2: &mut [f32],
5809    pool: Option<&Pool>,
5810) {
5811    q2tp_matvec2_mode(bytes, x1, x2, rows, cols, o1, o2, pool, false);
5812}
5813
5814#[allow(clippy::too_many_arguments)]
5815fn q2tp_affine_matvec2(
5816    bytes: &[u8],
5817    x1: &[f32],
5818    x2: &[f32],
5819    rows: usize,
5820    cols: usize,
5821    o1: &mut [f32],
5822    o2: &mut [f32],
5823    pool: Option<&Pool>,
5824) {
5825    q2tp_matvec2_mode(bytes, x1, x2, rows, cols, o1, o2, pool, true);
5826}
5827
5828#[allow(clippy::too_many_arguments)]
5829fn q2tp_matvec2_mode(
5830    bytes: &[u8],
5831    x1: &[f32],
5832    x2: &[f32],
5833    rows: usize,
5834    cols: usize,
5835    o1: &mut [f32],
5836    o2: &mut [f32],
5837    pool: Option<&Pool>,
5838    affine: bool,
5839) {
5840    let gpr = cols / GROUP_SIZE;
5841    let v = Q4tpView::new_q2(bytes, rows, cols);
5842    let (p1, p2) = (SendMut(o1.as_mut_ptr()), SendMut(o2.as_mut_ptr()));
5843    let run = |start: usize, end: usize| {
5844        let mut sc = vec![0f32; gpr];
5845        for r in start..end {
5846            v.scales_into(r, gpr, &mut sc);
5847            // SAFETY: disjoint row ranges per worker.
5848            unsafe {
5849                *p1.at(r) = if affine {
5850                    q2tp_affine_row_exact(v.nib, r, gpr, x1, &sc)
5851                } else {
5852                    q2tp_row_exact(v.nib, r, gpr, x1, &sc)
5853                };
5854                *p2.at(r) = if affine {
5855                    q2tp_affine_row_exact(v.nib, r, gpr, x2, &sc)
5856                } else {
5857                    q2tp_row_exact(v.nib, r, gpr, x2, &sc)
5858                };
5859            }
5860        }
5861    };
5862    dispatch_rows(pool, rows, &run);
5863}
5864
5865/// Batched q2tp matmat: scalar row kernel over every batch column. CPU
5866/// prefill only — decode rides the graph, so plain and correct beats
5867/// clever here.
5868/// Test doors into the host 2-bit kernels: the stand's heap corruption
5869/// pointed at down-shaped tensors, and the private fns need a way to be
5870/// held to a reference without a model file around them.
5871pub fn q2tp_matvec_for_test(bytes: &[u8], x: &[f32], rows: usize, cols: usize, out: &mut [f32]) {
5872    // The facade IS the reference: encoder oracles hold requant output
5873    // to the exact scalar walk. The production dispatch may take the i8
5874    // fast path, whose error scale is the ACTIVATIONS' — a different
5875    // claim than the encoder correctness these tests pin.
5876    let gpr = cols / GROUP_SIZE;
5877    let v = Q4tpView::new_q2(bytes, rows, cols);
5878    with_krow(gpr, |sc| {
5879        for r in 0..rows {
5880            v.scales_into(r, gpr, sc);
5881            out[r] = q2tp_row_exact(v.nib, r, gpr, x, sc);
5882        }
5883    });
5884}
5885
5886/// Test door for the descriptor-specific fused affine decode.  Production
5887/// callers select this through a validated Prism header, never by dtype alone.
5888pub fn q2tp_affine_matvec_for_test(
5889    bytes: &[u8],
5890    x: &[f32],
5891    rows: usize,
5892    cols: usize,
5893    out: &mut [f32],
5894) {
5895    q2tp_affine_matvec(bytes, x, rows, cols, out, None);
5896}
5897
5898pub fn q2tp_matmat_for_test(
5899    bytes: &[u8],
5900    xs_all: &[f32],
5901    b: usize,
5902    rows: usize,
5903    cols: usize,
5904    out: &mut [f32],
5905) {
5906    q2tp_matmat(bytes, xs_all, b, rows, cols, out, None);
5907}
5908
5909fn q2tp_matmat(
5910    bytes: &[u8],
5911    xs_all: &[f32],
5912    b: usize,
5913    rows: usize,
5914    cols: usize,
5915    out: &mut [f32],
5916    pool: Option<&Pool>,
5917) {
5918    q2tp_matmat_mode(bytes, xs_all, b, rows, cols, out, pool, false);
5919}
5920
5921fn q2tp_affine_matmat(
5922    bytes: &[u8],
5923    xs_all: &[f32],
5924    b: usize,
5925    rows: usize,
5926    cols: usize,
5927    out: &mut [f32],
5928    pool: Option<&Pool>,
5929) {
5930    q2tp_matmat_mode(bytes, xs_all, b, rows, cols, out, pool, true);
5931}
5932
5933fn q2tp_matmat_mode(
5934    bytes: &[u8],
5935    xs_all: &[f32],
5936    b: usize,
5937    rows: usize,
5938    cols: usize,
5939    out: &mut [f32],
5940    pool: Option<&Pool>,
5941    affine: bool,
5942) {
5943    debug_assert_eq!(out.len(), b * rows);
5944    let gpr = cols / GROUP_SIZE;
5945    let v = Q4tpView::new_q2(bytes, rows, cols);
5946    let out_addr = SendMut(out.as_mut_ptr());
5947    let run = |start: usize, end: usize| {
5948        let mut sc = vec![0f32; gpr];
5949        for r in start..end {
5950            v.scales_into(r, gpr, &mut sc);
5951            for bi in 0..b {
5952                let x = &xs_all[bi * cols..(bi + 1) * cols];
5953                // SAFETY: disjoint row ranges per worker.
5954                unsafe {
5955                    *out_addr.at(bi * rows + r) = if affine {
5956                        q2tp_affine_row_exact(v.nib, r, gpr, x, &sc)
5957                    } else {
5958                        q2tp_row_exact(v.nib, r, gpr, x, &sc)
5959                    }
5960                };
5961            }
5962        }
5963    };
5964    dispatch_rows(pool, rows, &run);
5965}
5966
5967/// The pre-vectorised shape, kept for A/B (`CMF_Q4TP_V1=1`): the
5968/// horizontal add lands once per group per column instead of once per
5969/// row. Same weights, same activations — only the reduction differs.
5970#[cfg(target_arch = "aarch64")]
5971#[target_feature(enable = "neon,dotprod")]
5972unsafe fn dot_q4tp_row_1x4_sdot_v1(
5973    nib: &[u8],
5974    r: usize,
5975    gpr: usize,
5976    xs: [&[i8]; 4],
5977    scales: &[f32],
5978) -> [f32; 4] {
5979    unsafe {
5980        use core::arch::aarch64::*;
5981        use core::arch::asm;
5982        let lomask = vdupq_n_u8(0x0F);
5983        let eight = vdupq_n_s8(8);
5984        let (mut f0, mut f1, mut f2, mut f3) = (0f32, 0f32, 0f32, 0f32);
5985        for gi in 0..gpr {
5986            let t = nib.as_ptr().add((r * gpr + gi) * Q4TP_NIB);
5987            let s = *scales.get_unchecked(gi);
5988            let bb = vld1q_u8(t);
5989            let lo = vandq_u8(bb, lomask);
5990            let hi = vshrq_n_u8::<4>(bb);
5991            let e0 = vsubq_s8(vreinterpretq_s8_u8(vzip1q_u8(lo, hi)), eight);
5992            let e1 = vsubq_s8(vreinterpretq_s8_u8(vzip2q_u8(lo, hi)), eight);
5993            let mut d = [0f32; 4];
5994            for (k, dk) in d.iter_mut().enumerate() {
5995                let x0 = vld1q_s8(xs[k].as_ptr().add(gi * GROUP_SIZE));
5996                let x1 = vld1q_s8(xs[k].as_ptr().add(gi * GROUP_SIZE + 16));
5997                let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
5998                asm!(
5999                    "sdot {a0:v}.4s, {e0:v}.16b, {x0:v}.16b",
6000                    "sdot {a1:v}.4s, {e1:v}.16b, {x1:v}.16b",
6001                    a0 = inout(vreg) a0, a1 = inout(vreg) a1,
6002                    e0 = in(vreg) e0, x0 = in(vreg) x0, e1 = in(vreg) e1, x1 = in(vreg) x1,
6003                    options(pure, nomem, nostack),
6004                );
6005                *dk = vaddvq_s32(vaddq_s32(a0, a1)) as f32 * s;
6006            }
6007            f0 += d[0];
6008            f1 += d[1];
6009            f2 += d[2];
6010            f3 += d[3];
6011        }
6012        [f0, f1, f2, f3]
6013    }
6014}
6015
6016/// Which q4tp batch kernel to run: 1 = the previous one, 2 = the tuned
6017/// one, 0 = decide from the CPU. An atomic rather than a `OnceLock` so a
6018/// benchmark can alternate the two inside one process, where the machine's
6019/// mood — a shared box drifts ±25% between runs — is the same for both.
6020/// What the two mean is per-architecture: on x86 the blocked AVX-512 path
6021/// against the per-column one, on ARM the two reduction shapes.
6022#[allow(dead_code)]
6023static Q4TP_ALT: std::sync::atomic::AtomicU8 = std::sync::atomic::AtomicU8::new(0);
6024
6025/// Blocking pays on x86 only with 512-bit VNNI. With AVX2 alone, four
6026/// columns sharing an unpack still measured slower than the per-column
6027/// path (23.2 ms against 19.4 on a 48-thread EPYC), because that path
6028/// already dequantizes the row once — so the blocked kernel bought a
6029/// second unpack-free pass at the price of half the vector width.
6030#[cfg(target_arch = "x86_64")]
6031fn q4tp_blocked_x86() -> bool {
6032    match Q4TP_ALT.load(std::sync::atomic::Ordering::Relaxed) {
6033        1 => false,
6034        // A forced ON still asks the CPU. The switch exists so a bench can
6035        // pick a kernel, not so it can promise instructions the machine
6036        // does not have — CI caught that as a SIGILL on a runner without
6037        // AVX-512, where the parity test had turned the path on by hand.
6038        2 => avx512vnni_enabled(),
6039        // Deliberately not cached back into the switch: both gates below
6040        // hold their own `OnceLock`, and latching their answer here would
6041        // make a test's override outlive the test that set it.
6042        _ => blocked_enabled() && avx512vnni_enabled(),
6043    }
6044}
6045
6046/// `CMF_Q4TP_V1=1` picks the old reduction shape (A/B only).
6047#[cfg(target_arch = "aarch64")]
6048#[allow(dead_code)]
6049fn q4tp_v1() -> bool {
6050    match Q4TP_ALT.load(std::sync::atomic::Ordering::Relaxed) {
6051        1 => true,
6052        2 => false,
6053        _ => {
6054            static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
6055            *ON.get_or_init(|| std::env::var("CMF_Q4TP_V1").is_ok_and(|v| v != "0"))
6056        }
6057    }
6058}
6059
6060/// Two weight rows against eight columns. The activation load is the
6061/// same for both rows, so it is paid once for twice the arithmetic, and
6062/// sixteen accumulator chains run where eight did — which is what a kernel
6063/// retiring 0.29 instructions a cycle is short of. Register pressure is
6064/// the limit: sixteen `zmm` accumulators, two weight tiles, one
6065/// activation, of thirty-two.
6066///
6067/// Four rows by four columns spends the same sixteen accumulators the
6068/// other way and measured worse — 1488 GFLOP/s against 1644 — so the
6069/// unpack, which four rows pay twice as often, costs more than the extra
6070/// sharing of one activation load buys.
6071#[cfg(target_arch = "x86_64")]
6072#[target_feature(enable = "avx512f,avx512bw,avx512vnni")]
6073unsafe fn dot_q4tp_2x8_avx512(
6074    nib: &[u8],
6075    r0: usize,
6076    gpr: usize,
6077    xs: [&[i8]; 8],
6078    sc0: &[f32],
6079    sc1: &[f32],
6080) -> [[f32; 8]; 2] {
6081    // SAFETY: as dot_q4tp_row_1x8_avx512, two adjacent rows at once; the
6082    // caller guarantees r0 + 1 < rows and the ISA.
6083    unsafe {
6084        use core::arch::x86_64::*;
6085        let lomask = _mm256_set1_epi8(0x0F);
6086        let eight = _mm256_set1_epi8(8);
6087        let zero = _mm512_setzero_si512();
6088        let mut v0 = [_mm512_setzero_ps(); 8];
6089        let mut v1 = [_mm512_setzero_ps(); 8];
6090        let pairs = gpr / 2;
6091        let unpack = |r: usize, gi: usize| -> (__m512i, __mmask64) {
6092            let t = nib.as_ptr().add((r * gpr + gi) * Q4TP_NIB);
6093            let bb = _mm256_loadu_si256(t as *const __m256i);
6094            let lo = _mm256_and_si256(bb, lomask);
6095            let hi = _mm256_and_si256(_mm256_srli_epi16::<4>(bb), lomask);
6096            let ul = _mm256_sub_epi8(_mm256_unpacklo_epi8(lo, hi), eight);
6097            let uh = _mm256_sub_epi8(_mm256_unpackhi_epi8(lo, hi), eight);
6098            let cat = _mm512_inserti64x4::<1>(_mm512_castsi256_si512(ul), uh);
6099            let w = _mm512_shuffle_i64x2::<0b11_01_10_00>(cat, cat);
6100            (_mm512_abs_epi8(w), _mm512_movepi8_mask(w))
6101        };
6102        for gp in 0..pairs {
6103            let gi = gp * 2;
6104            let (wa0, neg0) = unpack(r0, gi);
6105            let (wa1, neg1) = unpack(r0 + 1, gi);
6106            let off = gi * GROUP_SIZE;
6107            let sv = |sc: &[f32]| {
6108                _mm512_insertf32x8::<1>(
6109                    _mm512_castps256_ps512(_mm256_set1_ps(*sc.get_unchecked(gi))),
6110                    _mm256_set1_ps(*sc.get_unchecked(gi + 1)),
6111                )
6112            };
6113            let s0 = sv(sc0);
6114            let s1 = sv(sc1);
6115            for k in 0..8 {
6116                let xv = _mm512_loadu_si512(xs[k].as_ptr().add(off) as *const __m512i);
6117                let d0 = _mm512_cvtepi32_ps(_mm512_dpbusd_epi32(
6118                    zero,
6119                    wa0,
6120                    _mm512_mask_sub_epi8(xv, neg0, zero, xv),
6121                ));
6122                let d1 = _mm512_cvtepi32_ps(_mm512_dpbusd_epi32(
6123                    zero,
6124                    wa1,
6125                    _mm512_mask_sub_epi8(xv, neg1, zero, xv),
6126                ));
6127                v0[k] = _mm512_fmadd_ps(d0, s0, v0[k]);
6128                v1[k] = _mm512_fmadd_ps(d1, s1, v1[k]);
6129            }
6130        }
6131        let mut acc = [[0f32; 8]; 2];
6132        for k in 0..8 {
6133            acc[0][k] = _mm512_reduce_add_ps(v0[k]);
6134            acc[1][k] = _mm512_reduce_add_ps(v1[k]);
6135        }
6136        if gpr % 2 == 1 {
6137            let off = (gpr - 1) * GROUP_SIZE;
6138            for j in off..off + GROUP_SIZE {
6139                let (w0, sa) = q4tp_outlier(nib, r0, gpr, j, sc0);
6140                let (w1, sb) = q4tp_outlier(nib, r0 + 1, gpr, j, sc1);
6141                for k in 0..8 {
6142                    let x = *xs[k].get_unchecked(j) as f32;
6143                    acc[0][k] += w0 * sa * x;
6144                    acc[1][k] += w1 * sb * x;
6145                }
6146            }
6147        }
6148        acc
6149    }
6150}
6151
6152/// The same, eight columns at a time. One unpack then feeds twice as many
6153/// activation streams, so a wide batch reads the weight tile half as
6154/// often; the price is eight accumulators live at once. Measured 9.0 ->
6155/// 8.3 ms at 9216x2304, b=296 on a 48-thread EPYC 9B45.
6156#[cfg(target_arch = "x86_64")]
6157#[target_feature(enable = "avx512f,avx512bw,avx512vnni")]
6158unsafe fn dot_q4tp_row_1x8_avx512(
6159    nib: &[u8],
6160    r: usize,
6161    gpr: usize,
6162    xs: [&[i8]; 8],
6163    scales: &[f32],
6164) -> [f32; 8] {
6165    // SAFETY: as dot_q4tp_row_1x4_avx2; caller guarantees the ISA.
6166    unsafe {
6167        use core::arch::x86_64::*;
6168        let lomask = _mm256_set1_epi8(0x0F);
6169        let eight = _mm256_set1_epi8(8);
6170        let zero = _mm512_setzero_si512();
6171        let (mut v0, mut v1, mut v2, mut v3) = (
6172            _mm512_setzero_ps(),
6173            _mm512_setzero_ps(),
6174            _mm512_setzero_ps(),
6175            _mm512_setzero_ps(),
6176        );
6177        let (mut v4, mut v5, mut v6, mut v7) = (
6178            _mm512_setzero_ps(),
6179            _mm512_setzero_ps(),
6180            _mm512_setzero_ps(),
6181            _mm512_setzero_ps(),
6182        );
6183        let pairs = gpr / 2;
6184        for gp in 0..pairs {
6185            let gi = gp * 2;
6186            let t = nib.as_ptr().add((r * gpr + gi) * Q4TP_NIB);
6187            let bb = _mm256_loadu_si256(t as *const __m256i);
6188            let lo = _mm256_and_si256(bb, lomask);
6189            let hi = _mm256_and_si256(_mm256_srli_epi16::<4>(bb), lomask);
6190            // `unpack` works per 128-bit lane, so the halves come out as
6191            // [A.lo, B.lo] and [A.hi, B.hi]; the shuffle reorders the four
6192            // 128-bit lanes into the weights' natural order, which is what
6193            // the straight activation load expects.
6194            let ul = _mm256_sub_epi8(_mm256_unpacklo_epi8(lo, hi), eight);
6195            let uh = _mm256_sub_epi8(_mm256_unpackhi_epi8(lo, hi), eight);
6196            let cat = _mm512_inserti64x4::<1>(_mm512_castsi256_si512(ul), uh);
6197            let w = _mm512_shuffle_i64x2::<0b11_01_10_00>(cat, cat);
6198            let wabs = _mm512_abs_epi8(w);
6199            let neg = _mm512_movepi8_mask(w);
6200            let off = gi * GROUP_SIZE;
6201            let sv = _mm512_insertf32x8::<1>(
6202                _mm512_castps256_ps512(_mm256_set1_ps(*scales.get_unchecked(gi))),
6203                _mm256_set1_ps(*scales.get_unchecked(gi + 1)),
6204            );
6205            let dot = |x: &[i8]| -> __m512 {
6206                let xv = _mm512_loadu_si512(x.as_ptr().add(off) as *const __m512i);
6207                let sx = _mm512_mask_sub_epi8(xv, neg, zero, xv);
6208                _mm512_cvtepi32_ps(_mm512_dpbusd_epi32(zero, wabs, sx))
6209            };
6210            v0 = _mm512_fmadd_ps(dot(xs[0]), sv, v0);
6211            v1 = _mm512_fmadd_ps(dot(xs[1]), sv, v1);
6212            v2 = _mm512_fmadd_ps(dot(xs[2]), sv, v2);
6213            v3 = _mm512_fmadd_ps(dot(xs[3]), sv, v3);
6214            v4 = _mm512_fmadd_ps(dot(xs[4]), sv, v4);
6215            v5 = _mm512_fmadd_ps(dot(xs[5]), sv, v5);
6216            v6 = _mm512_fmadd_ps(dot(xs[6]), sv, v6);
6217            v7 = _mm512_fmadd_ps(dot(xs[7]), sv, v7);
6218        }
6219        let mut acc = [
6220            _mm512_reduce_add_ps(v0),
6221            _mm512_reduce_add_ps(v1),
6222            _mm512_reduce_add_ps(v2),
6223            _mm512_reduce_add_ps(v3),
6224            _mm512_reduce_add_ps(v4),
6225            _mm512_reduce_add_ps(v5),
6226            _mm512_reduce_add_ps(v6),
6227            _mm512_reduce_add_ps(v7),
6228        ];
6229        // An odd group count leaves one group over; the narrow kernel
6230        // finishes it rather than the tail being a special case here.
6231        if gpr % 2 == 1 {
6232            let off = (gpr - 1) * GROUP_SIZE;
6233            for j in off..off + GROUP_SIZE {
6234                let (w, s) = q4tp_outlier(nib, r, gpr, j, scales);
6235                let ws = w * s;
6236                for k in 0..8 {
6237                    acc[k] += ws * *xs[k].get_unchecked(j) as f32;
6238                }
6239            }
6240        }
6241        acc
6242    }
6243}
6244
6245/// The same four columns, 512 bits wide. Two groups (64 weights) ride one
6246/// unpack and one `vpdpbusd`, where AVX2 needs two unpacks and four
6247/// `maddubs`/`madd` pairs — about 2.3x fewer instructions for the same
6248/// arithmetic. The two groups carry different scales, so the fma takes a
6249/// vector whose halves hold each group's scale rather than a broadcast.
6250///
6251/// There is no 512-bit `vpsignb`, so the activation's sign is applied by
6252/// negating under a mask taken from the weight's sign bits. That mask is
6253/// per-tile, so it is hoisted out of the column loop and the per-column
6254/// cost stays exactly one instruction, as with `sign_epi8`. Weights of
6255/// zero are not zeroed by the mask trick and do not need to be: their
6256/// magnitude is zero, so the product is.
6257#[cfg(target_arch = "x86_64")]
6258#[target_feature(enable = "avx512f,avx512bw,avx512vnni")]
6259unsafe fn dot_q4tp_row_1x4_avx512(
6260    nib: &[u8],
6261    r: usize,
6262    gpr: usize,
6263    xs: [&[i8]; 4],
6264    scales: &[f32],
6265) -> [f32; 4] {
6266    // SAFETY: as dot_q4tp_row_1x4_avx2; caller guarantees the ISA.
6267    unsafe {
6268        use core::arch::x86_64::*;
6269        let lomask = _mm256_set1_epi8(0x0F);
6270        let eight = _mm256_set1_epi8(8);
6271        let zero = _mm512_setzero_si512();
6272        let (mut v0, mut v1, mut v2, mut v3) = (
6273            _mm512_setzero_ps(),
6274            _mm512_setzero_ps(),
6275            _mm512_setzero_ps(),
6276            _mm512_setzero_ps(),
6277        );
6278        let pairs = gpr / 2;
6279        for gp in 0..pairs {
6280            let gi = gp * 2;
6281            let t = nib.as_ptr().add((r * gpr + gi) * Q4TP_NIB);
6282            let bb = _mm256_loadu_si256(t as *const __m256i);
6283            let lo = _mm256_and_si256(bb, lomask);
6284            let hi = _mm256_and_si256(_mm256_srli_epi16::<4>(bb), lomask);
6285            // `unpack` works per 128-bit lane, so the halves come out as
6286            // [A.lo, B.lo] and [A.hi, B.hi]; the shuffle reorders the four
6287            // 128-bit lanes into the weights' natural order, which is what
6288            // the straight activation load expects.
6289            let ul = _mm256_sub_epi8(_mm256_unpacklo_epi8(lo, hi), eight);
6290            let uh = _mm256_sub_epi8(_mm256_unpackhi_epi8(lo, hi), eight);
6291            let cat = _mm512_inserti64x4::<1>(_mm512_castsi256_si512(ul), uh);
6292            let w = _mm512_shuffle_i64x2::<0b11_01_10_00>(cat, cat);
6293            let wabs = _mm512_abs_epi8(w);
6294            let neg = _mm512_movepi8_mask(w);
6295            let off = gi * GROUP_SIZE;
6296            let sv = _mm512_insertf32x8::<1>(
6297                _mm512_castps256_ps512(_mm256_set1_ps(*scales.get_unchecked(gi))),
6298                _mm256_set1_ps(*scales.get_unchecked(gi + 1)),
6299            );
6300            let dot = |x: &[i8]| -> __m512 {
6301                let xv = _mm512_loadu_si512(x.as_ptr().add(off) as *const __m512i);
6302                let sx = _mm512_mask_sub_epi8(xv, neg, zero, xv);
6303                _mm512_cvtepi32_ps(_mm512_dpbusd_epi32(zero, wabs, sx))
6304            };
6305            v0 = _mm512_fmadd_ps(dot(xs[0]), sv, v0);
6306            v1 = _mm512_fmadd_ps(dot(xs[1]), sv, v1);
6307            v2 = _mm512_fmadd_ps(dot(xs[2]), sv, v2);
6308            v3 = _mm512_fmadd_ps(dot(xs[3]), sv, v3);
6309        }
6310        let mut acc = [
6311            _mm512_reduce_add_ps(v0),
6312            _mm512_reduce_add_ps(v1),
6313            _mm512_reduce_add_ps(v2),
6314            _mm512_reduce_add_ps(v3),
6315        ];
6316        // An odd group count leaves one group over; the narrow kernel
6317        // finishes it rather than the tail being a special case here.
6318        if gpr % 2 == 1 {
6319            let off = (gpr - 1) * GROUP_SIZE;
6320            for j in off..off + GROUP_SIZE {
6321                let (w, s) = q4tp_outlier(nib, r, gpr, j, scales);
6322                let ws = w * s;
6323                for k in 0..4 {
6324                    acc[k] += ws * *xs[k].get_unchecked(j) as f32;
6325                }
6326            }
6327        }
6328        acc
6329    }
6330}
6331
6332/// Four batch columns against one q4tp row: the tile is unpacked ONCE and
6333/// spent on four activation streams, which is where a prefill batch stops
6334/// being weight-bandwidth-bound. Twin of `dot_q4t_row_1x4_sdot`.
6335#[cfg(target_arch = "aarch64")]
6336#[target_feature(enable = "neon,dotprod")]
6337unsafe fn dot_q4tp_row_1x4_sdot(
6338    nib: &[u8],
6339    r: usize,
6340    gpr: usize,
6341    xs: [&[i8]; 4],
6342    scales: &[f32],
6343) -> [f32; 4] {
6344    // SAFETY: see dot_q4tp_row_sdot; every xs[k] is gpr·GROUP_SIZE long.
6345    unsafe {
6346        use core::arch::aarch64::*;
6347        use core::arch::asm;
6348        let lomask = vdupq_n_u8(0x0F);
6349        let eight = vdupq_n_s8(8);
6350        // Named accumulators, NOT an array indexed by a loop variable: the
6351        // latter does not stay in registers (the same defect cost 2x in the
6352        // AVX2 q4t kernel and again in WGSL).
6353        //
6354        // They are VECTORS, and the horizontal add happens once at the end
6355        // instead of once per group per column. `vaddvq` is a cross-lane
6356        // reduction — with 72 groups and four columns the old shape paid
6357        // 288 of them per row, each one a dependency stall the pipeline
6358        // cannot hide, to save four float adds. The group's scale now
6359        // rides an fma into the lane accumulators, so the arithmetic per
6360        // group is one convert and one fma. Summation order changes (the
6361        // lanes carry independent partial sums), which is the same
6362        // round-off class the SDOT path already lives in — the strict
6363        // kernel (`CMF_SDOT=0`, what `cortiq ppl` runs) is unchanged and
6364        // stays the reference.
6365        let (mut v0, mut v1, mut v2, mut v3) = (
6366            vdupq_n_f32(0.0),
6367            vdupq_n_f32(0.0),
6368            vdupq_n_f32(0.0),
6369            vdupq_n_f32(0.0),
6370        );
6371        for gi in 0..gpr {
6372            let t = nib.as_ptr().add((r * gpr + gi) * Q4TP_NIB);
6373            let s = *scales.get_unchecked(gi);
6374            let bb = vld1q_u8(t);
6375            let lo = vandq_u8(bb, lomask);
6376            let hi = vshrq_n_u8::<4>(bb);
6377            let e0 = vsubq_s8(vreinterpretq_s8_u8(vzip1q_u8(lo, hi)), eight);
6378            let e1 = vsubq_s8(vreinterpretq_s8_u8(vzip2q_u8(lo, hi)), eight);
6379            let off = gi * GROUP_SIZE;
6380            let dot4 = |x: &[i8]| -> int32x4_t {
6381                let x0 = vld1q_s8(x.as_ptr().add(off));
6382                let x1 = vld1q_s8(x.as_ptr().add(off + 16));
6383                let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
6384                asm!(
6385                    "sdot {a0:v}.4s, {e0:v}.16b, {x0:v}.16b",
6386                    "sdot {a1:v}.4s, {e1:v}.16b, {x1:v}.16b",
6387                    a0 = inout(vreg) a0, a1 = inout(vreg) a1,
6388                    e0 = in(vreg) e0, x0 = in(vreg) x0, e1 = in(vreg) e1, x1 = in(vreg) x1,
6389                    options(pure, nomem, nostack),
6390                );
6391                vaddq_s32(a0, a1)
6392            };
6393            v0 = vfmaq_n_f32(v0, vcvtq_f32_s32(dot4(xs[0])), s);
6394            v1 = vfmaq_n_f32(v1, vcvtq_f32_s32(dot4(xs[1])), s);
6395            v2 = vfmaq_n_f32(v2, vcvtq_f32_s32(dot4(xs[2])), s);
6396            v3 = vfmaq_n_f32(v3, vcvtq_f32_s32(dot4(xs[3])), s);
6397        }
6398        [
6399            vaddvq_f32(v0),
6400            vaddvq_f32(v1),
6401            vaddvq_f32(v2),
6402            vaddvq_f32(v3),
6403        ]
6404    }
6405}
6406
6407/// Fused q4tp matmat — the same three arms `q4t_matmat` has. Shipping only
6408/// the scalar one made Nanbeige-3B decode at 1.2 tok/s against q4t's 5.9:
6409/// the format was fine, the missing arms were the whole regression.
6410fn q4tp_matmat(
6411    bytes: &[u8],
6412    xs_all: &[f32],
6413    b: usize,
6414    rows: usize,
6415    cols: usize,
6416    out: &mut [f32],
6417    pool: Option<&Pool>,
6418) {
6419    debug_assert_eq!(out.len(), b * rows);
6420    let gpr = cols / GROUP_SIZE;
6421    let v = Q4tpView::new(bytes, rows, cols);
6422
6423    // Wide batches ride the AMX through a dequant-tile sgemm, as in q4t.
6424    #[cfg(target_os = "macos")]
6425    if b >= 8 && rows * cols >= 500_000 && accel_gemm_enabled() {
6426        dequant_matmat_accel(
6427            &|r, dst| {
6428                let mut sc = [0f32; 32];
6429                let mut scv;
6430                let s: &[f32] = if gpr <= 32 {
6431                    v.scales_into(r, gpr, &mut sc);
6432                    &sc[..gpr]
6433                } else {
6434                    scv = vec![0f32; gpr];
6435                    v.scales_into(r, gpr, &mut scv);
6436                    &scv
6437                };
6438                for gi in 0..gpr {
6439                    let tile = &v.nib[(r * gpr + gi) * Q4TP_NIB..(r * gpr + gi + 1) * Q4TP_NIB];
6440                    for (k, &bb) in tile.iter().enumerate() {
6441                        dst[gi * GROUP_SIZE + k * 2] = ((bb & 0x0F) as f32 - 8.0) * s[gi];
6442                        dst[gi * GROUP_SIZE + k * 2 + 1] =
6443                            (((bb >> 4) & 0x0F) as f32 - 8.0) * s[gi];
6444                    }
6445                }
6446            },
6447            xs_all,
6448            b,
6449            rows,
6450            cols,
6451            out,
6452            pool,
6453        );
6454        return;
6455    }
6456
6457    let out_addr = SendMut(out.as_mut_ptr());
6458    if a8w8_enabled() {
6459        let acts: Vec<SplitAct> = (0..b)
6460            .map(|bi| split_act(&xs_all[bi * cols..(bi + 1) * cols]))
6461            .collect();
6462        let acts = &acts;
6463        #[cfg(target_arch = "aarch64")]
6464        let blocked_ok = sdot_enabled() && blocked_enabled();
6465        // x86 gets the same blocking: one tile unpack spent on four
6466        // columns. Without it every column re-decoded the row, which is
6467        // why a 48-core EPYC measured a sixth of an M4's per-core rate.
6468        // The gate is `avx2_enabled`, as in q4t — `sdot_enabled` answers
6469        // for ARM's dotprod and is hard-wired false everywhere else, so
6470        // asking it here left the whole blocked path unreachable on x86.
6471        #[cfg(target_arch = "x86_64")]
6472        let blocked_ok = q4tp_blocked_x86() && !row_exact();
6473        #[cfg(not(any(target_arch = "aarch64", target_arch = "x86_64")))]
6474        let blocked_ok = false;
6475        // Columns are swept in panels that fit L2. Without this a
6476        // row-pair walks every activation in the batch — 4.8 MB at
6477        // 512x512 — and does it again for the next pair, so the whole
6478        // batch streams out of the shared cache once per row. Measured
6479        // 800 GB/s of it, flat across batch sizes, which is the signature
6480        // of a loop bound by traffic rather than by arithmetic. A panel of
6481        // 256 columns is 590 KB beside 221 KB of this worker's weights:
6482        // both stay resident and the batch crosses L3 once instead of
6483        // once per row.
6484        let panel_cols: usize = std::env::var("CMF_Q4TP_PANEL")
6485            .ok()
6486            .and_then(|v| v.parse().ok())
6487            .filter(|v| *v > 0)
6488            .unwrap_or(256);
6489        let run = |start: usize, end: usize| {
6490            for abase in (0..acts.len()).step_by(panel_cols) {
6491                let alen = (acts.len() - abase).min(panel_cols);
6492                let mut sc = vec![0f32; gpr];
6493                #[cfg(target_arch = "x86_64")]
6494                let mut r_lo = start;
6495                #[cfg(target_arch = "x86_64")]
6496                if blocked_ok && alen >= 8 {
6497                    let mut sc1 = vec![0f32; gpr];
6498                    while r_lo + 2 <= end {
6499                        v.scales_into(r_lo, gpr, &mut sc);
6500                        v.scales_into(r_lo + 1, gpr, &mut sc1);
6501                        let mut bi = 0usize;
6502                        while bi + 8 <= alen {
6503                            let xs = [
6504                                acts[abase + bi].xq.as_slice(),
6505                                acts[abase + bi + 1].xq.as_slice(),
6506                                acts[abase + bi + 2].xq.as_slice(),
6507                                acts[abase + bi + 3].xq.as_slice(),
6508                                acts[abase + bi + 4].xq.as_slice(),
6509                                acts[abase + bi + 5].xq.as_slice(),
6510                                acts[abase + bi + 6].xq.as_slice(),
6511                                acts[abase + bi + 7].xq.as_slice(),
6512                            ];
6513                            let d = unsafe { dot_q4tp_2x8_avx512(v.nib, r_lo, gpr, xs, &sc, &sc1) };
6514                            for (row, dr, scr) in [(r_lo, &d[0], &sc), (r_lo + 1, &d[1], &sc1)] {
6515                                for k in 0..8 {
6516                                    let act = &acts[abase + bi + k];
6517                                    let mut acc = dr[k] * act.sx;
6518                                    for &(j, xv) in &act.outliers {
6519                                        let (w, s) = q4tp_outlier(v.nib, row, gpr, j, scr);
6520                                        acc += w * s * xv;
6521                                    }
6522                                    // SAFETY: disjoint (bi, r) cells per worker.
6523                                    unsafe { *out_addr.at((abase + bi + k) * rows + row) = acc };
6524                                }
6525                            }
6526                            bi += 8;
6527                        }
6528                        // Columns past the last group of eight, both rows —
6529                        // the same single-row kernel the tail below uses.
6530                        for row in [r_lo, r_lo + 1] {
6531                            let scr: &[f32] = if row == r_lo { &sc } else { &sc1 };
6532                            for b2 in bi..alen {
6533                                let act = &acts[abase + b2];
6534                                let xs4 = [
6535                                    act.xq.as_slice(),
6536                                    act.xq.as_slice(),
6537                                    act.xq.as_slice(),
6538                                    act.xq.as_slice(),
6539                                ];
6540                                let d =
6541                                    unsafe { dot_q4tp_row_1x4_avx512(v.nib, row, gpr, xs4, scr) };
6542                                let mut acc = d[0] * act.sx;
6543                                for &(j, xv) in &act.outliers {
6544                                    let (w, s) = q4tp_outlier(v.nib, row, gpr, j, scr);
6545                                    acc += w * s * xv;
6546                                }
6547                                // SAFETY: disjoint (bi, r) cells per worker.
6548                                unsafe { *out_addr.at((abase + b2) * rows + row) = acc };
6549                            }
6550                        }
6551                        r_lo += 2;
6552                    }
6553                }
6554                #[cfg(target_arch = "x86_64")]
6555                let row_start = r_lo;
6556                #[cfg(not(target_arch = "x86_64"))]
6557                let row_start = start;
6558                for r in row_start..end {
6559                    v.scales_into(r, gpr, &mut sc);
6560                    let mut bi = 0usize;
6561                    #[cfg(target_arch = "x86_64")]
6562                    if blocked_ok {
6563                        while bi + 8 <= alen {
6564                            let xs = [
6565                                acts[abase + bi].xq.as_slice(),
6566                                acts[abase + bi + 1].xq.as_slice(),
6567                                acts[abase + bi + 2].xq.as_slice(),
6568                                acts[abase + bi + 3].xq.as_slice(),
6569                                acts[abase + bi + 4].xq.as_slice(),
6570                                acts[abase + bi + 5].xq.as_slice(),
6571                                acts[abase + bi + 6].xq.as_slice(),
6572                                acts[abase + bi + 7].xq.as_slice(),
6573                            ];
6574                            let d = unsafe { dot_q4tp_row_1x8_avx512(v.nib, r, gpr, xs, &sc) };
6575                            for k in 0..8 {
6576                                let act = &acts[abase + bi + k];
6577                                let mut acc = d[k] * act.sx;
6578                                for &(j, xv) in &act.outliers {
6579                                    let (w, s) = q4tp_outlier(v.nib, r, gpr, j, &sc);
6580                                    acc += w * s * xv;
6581                                }
6582                                // SAFETY: disjoint (bi, r) cells per worker.
6583                                unsafe { *out_addr.at((abase + bi + k) * rows + r) = acc };
6584                            }
6585                            bi += 8;
6586                        }
6587                        while bi + 4 <= alen {
6588                            let xs = [
6589                                acts[abase + bi].xq.as_slice(),
6590                                acts[abase + bi + 1].xq.as_slice(),
6591                                acts[abase + bi + 2].xq.as_slice(),
6592                                acts[abase + bi + 3].xq.as_slice(),
6593                            ];
6594                            let d = unsafe { dot_q4tp_row_1x4_avx512(v.nib, r, gpr, xs, &sc) };
6595                            for k in 0..4 {
6596                                let act = &acts[abase + bi + k];
6597                                let mut acc = d[k] * act.sx;
6598                                for &(j, xv) in &act.outliers {
6599                                    let (w, s) = q4tp_outlier(v.nib, r, gpr, j, &sc);
6600                                    acc += w * s * xv;
6601                                }
6602                                // SAFETY: disjoint (bi, r) cells per worker.
6603                                unsafe { *out_addr.at((abase + bi + k) * rows + r) = acc };
6604                            }
6605                            bi += 4;
6606                        }
6607                    }
6608                    #[cfg(target_arch = "aarch64")]
6609                    if blocked_ok {
6610                        while bi + 4 <= alen {
6611                            let xs = [
6612                                acts[abase + bi].xq.as_slice(),
6613                                acts[abase + bi + 1].xq.as_slice(),
6614                                acts[abase + bi + 2].xq.as_slice(),
6615                                acts[abase + bi + 3].xq.as_slice(),
6616                            ];
6617                            let d = unsafe {
6618                                if q4tp_v1() {
6619                                    dot_q4tp_row_1x4_sdot_v1(v.nib, r, gpr, xs, &sc)
6620                                } else {
6621                                    dot_q4tp_row_1x4_sdot(v.nib, r, gpr, xs, &sc)
6622                                }
6623                            };
6624                            for k in 0..4 {
6625                                let act = &acts[abase + bi + k];
6626                                let mut acc = d[k] * act.sx;
6627                                for &(j, xv) in &act.outliers {
6628                                    let (w, s) = q4tp_outlier(v.nib, r, gpr, j, &sc);
6629                                    acc += w * s * xv;
6630                                }
6631                                // SAFETY: disjoint (bi, r) cells per worker.
6632                                unsafe { *out_addr.at((abase + bi + k) * rows + r) = acc };
6633                            }
6634                            bi += 4;
6635                        }
6636                    }
6637                    let _ = blocked_ok;
6638                    while bi < alen {
6639                        let act = &acts[abase + bi];
6640                        let mut acc = dot_q4tp_row_i8(v.nib, r, gpr, &act.xq, &sc) * act.sx;
6641                        for &(j, xv) in &act.outliers {
6642                            let (w, s) = q4tp_outlier(v.nib, r, gpr, j, &sc);
6643                            acc += w * s * xv;
6644                        }
6645                        // SAFETY: disjoint (bi, r) cells per worker range.
6646                        unsafe { *out_addr.at((abase + bi) * rows + r) = acc };
6647                        bi += 1;
6648                    }
6649                }
6650            }
6651        };
6652        dispatch_rows(pool, rows, &run);
6653        return;
6654    }
6655
6656    let run = |start: usize, end: usize| {
6657        let mut sc = vec![0f32; gpr];
6658        for r in start..end {
6659            v.scales_into(r, gpr, &mut sc);
6660            for bi in 0..b {
6661                let x = &xs_all[bi * cols..(bi + 1) * cols];
6662                // SAFETY: disjoint (bi, r) cells per worker range.
6663                unsafe { *out_addr.at(bi * rows + r) = q4tp_row_exact(v.nib, r, gpr, x, &sc) };
6664            }
6665        }
6666    };
6667    dispatch_rows(pool, rows, &run);
6668}
6669
6670/// Fused q4_tiled matvec (dispatch mirrors `q4matvec`).
6671fn q4t_matvec(
6672    bytes: &[u8],
6673    x: &[f32],
6674    rows: usize,
6675    cols: usize,
6676    out: &mut [f32],
6677    pool: Option<&Pool>,
6678) {
6679    debug_assert_eq!(out.len(), rows);
6680    let gpr = cols / GROUP_SIZE;
6681    let out_addr = SendMut(out.as_mut_ptr());
6682    if a8w8_enabled() {
6683        let act = split_act(x);
6684        let run = move |start: usize, end: usize| {
6685            for r in start..end {
6686                let mut acc = dot_q4t_row_i8(bytes, r, gpr, &act.xq) * act.sx;
6687                for &(j, xv) in &act.outliers {
6688                    let (w, s) = q4t_outlier(bytes, r, gpr, j);
6689                    acc += w * s * xv;
6690                }
6691                // SAFETY: disjoint row ranges per worker.
6692                unsafe { *out_addr.at(r) = acc };
6693            }
6694        };
6695        dispatch_rows(pool, rows, &run);
6696        return;
6697    }
6698    let run = move |start: usize, end: usize| {
6699        for r in start..end {
6700            // SAFETY: disjoint row ranges per worker.
6701            unsafe { *out_addr.at(r) = q4t_row_exact(bytes, r, gpr, x) };
6702        }
6703    };
6704    dispatch_rows(pool, rows, &run);
6705}
6706
6707/// Fused two-input q4_tiled matvec (weights read once per pair).
6708#[allow(clippy::too_many_arguments)]
6709fn q4t_matvec2(
6710    bytes: &[u8],
6711    x1: &[f32],
6712    x2: &[f32],
6713    rows: usize,
6714    cols: usize,
6715    o1: &mut [f32],
6716    o2: &mut [f32],
6717    pool: Option<&Pool>,
6718) {
6719    let gpr = cols / GROUP_SIZE;
6720    let p1 = SendMut(o1.as_mut_ptr());
6721    let p2 = SendMut(o2.as_mut_ptr());
6722    if a8w8_enabled() {
6723        let a1 = split_act(x1);
6724        let a2 = split_act(x2);
6725        let run = move |start: usize, end: usize| {
6726            for r in start..end {
6727                let mut v1 = dot_q4t_row_i8(bytes, r, gpr, &a1.xq) * a1.sx;
6728                let mut v2 = dot_q4t_row_i8(bytes, r, gpr, &a2.xq) * a2.sx;
6729                for &(j, xv) in &a1.outliers {
6730                    let (w, s) = q4t_outlier(bytes, r, gpr, j);
6731                    v1 += w * s * xv;
6732                }
6733                for &(j, xv) in &a2.outliers {
6734                    let (w, s) = q4t_outlier(bytes, r, gpr, j);
6735                    v2 += w * s * xv;
6736                }
6737                // SAFETY: disjoint row ranges per worker.
6738                unsafe {
6739                    *p1.at(r) = v1;
6740                    *p2.at(r) = v2;
6741                }
6742            }
6743        };
6744        dispatch_rows(pool, rows, &run);
6745        return;
6746    }
6747    let run = move |start: usize, end: usize| {
6748        for r in start..end {
6749            // SAFETY: disjoint row ranges per worker.
6750            unsafe {
6751                *p1.at(r) = q4t_row_exact(bytes, r, gpr, x1);
6752                *p2.at(r) = q4t_row_exact(bytes, r, gpr, x2);
6753            }
6754        }
6755    };
6756    dispatch_rows(pool, rows, &run);
6757}
6758
6759/// Batched q4_tiled matmat: each row's tiles stream once per microbatch.
6760#[allow(clippy::too_many_arguments)]
6761/// Prefill GEMM through Accelerate for group-quantized codecs: a
6762/// caller-supplied row dequantizer fills f32 tiles (pool-parallel) and
6763/// each tile rides the AMX with one sgemm — the generic sibling of
6764/// `qmatmat_accel` (q8). Numerics are f32-GEMM (tolerance class);
6765/// decode (b=1) never takes this path.
6766#[cfg(target_os = "macos")]
6767fn dequant_matmat_accel(
6768    dequant_row: &(dyn Fn(usize, &mut [f32]) + Sync),
6769    xs_all: &[f32],
6770    b: usize,
6771    rows: usize,
6772    cols: usize,
6773    out: &mut [f32],
6774    pool: Option<&Pool>,
6775) {
6776    const TR: usize = 2048;
6777    thread_local! {
6778        static WTILE: std::cell::RefCell<Vec<f32>> = const { std::cell::RefCell::new(Vec::new()) };
6779    }
6780    WTILE.with(|wt| {
6781        let mut wtile = wt.borrow_mut();
6782        wtile.resize(TR * cols, 0.0);
6783        let mut r0 = 0usize;
6784        while r0 < rows {
6785            let tr = TR.min(rows - r0);
6786            let wt_addr = SendMut(wtile.as_mut_ptr());
6787            let run = |start: usize, end: usize| {
6788                for r in start..end {
6789                    // SAFETY: workers cover disjoint r ranges.
6790                    let dst = unsafe { std::slice::from_raw_parts_mut(wt_addr.at(r * cols), cols) };
6791                    dequant_row(r0 + r, dst);
6792                }
6793            };
6794            dispatch_rows(pool, tr, &run);
6795            unsafe {
6796                accel_blas::cblas_sgemm(
6797                    101, // RowMajor
6798                    111, // NoTrans A
6799                    112, // Trans B
6800                    b as i32,
6801                    tr as i32,
6802                    cols as i32,
6803                    1.0,
6804                    xs_all.as_ptr(),
6805                    cols as i32,
6806                    wtile.as_ptr(),
6807                    cols as i32,
6808                    0.0,
6809                    out.as_mut_ptr().add(r0),
6810                    rows as i32,
6811                );
6812            }
6813            r0 += tr;
6814        }
6815    });
6816}
6817
6818fn q4t_matmat(
6819    bytes: &[u8],
6820    xs_all: &[f32],
6821    b: usize,
6822    rows: usize,
6823    cols: usize,
6824    out: &mut [f32],
6825    pool: Option<&Pool>,
6826) {
6827    debug_assert_eq!(out.len(), b * rows);
6828    let gpr = cols / GROUP_SIZE;
6829    // Wide batches ride the AMX like q8's qmatmat: on Apple silicon
6830    // the dequant-tile sgemm is an order above the SDOT row loop for
6831    // prefill shapes (imagegen DiT forwards are exactly this).
6832    #[cfg(target_os = "macos")]
6833    if b >= 8 && rows * cols >= 500_000 && accel_gemm_enabled() {
6834        dequant_matmat_accel(
6835            &|r, dst| {
6836                for gi in 0..gpr {
6837                    let tile = &bytes[(r * gpr + gi) * Q4_TILE..(r * gpr + gi + 1) * Q4_TILE];
6838                    let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
6839                    for (k, &bb) in tile[2..].iter().enumerate() {
6840                        dst[gi * GROUP_SIZE + k * 2] = ((bb & 0x0F) as f32 - 8.0) * s;
6841                        dst[gi * GROUP_SIZE + k * 2 + 1] = (((bb >> 4) & 0x0F) as f32 - 8.0) * s;
6842                    }
6843                }
6844            },
6845            xs_all,
6846            b,
6847            rows,
6848            cols,
6849            out,
6850            pool,
6851        );
6852        return;
6853    }
6854    let out_addr = SendMut(out.as_mut_ptr());
6855    if a8w8_enabled() {
6856        let acts: Vec<SplitAct> = (0..b)
6857            .map(|bi| split_act(&xs_all[bi * cols..(bi + 1) * cols]))
6858            .collect();
6859        let acts = &acts;
6860        #[cfg(target_arch = "x86_64")]
6861        let blocked_ok = avx2_enabled() && blocked_enabled();
6862        #[cfg(target_arch = "aarch64")]
6863        let blocked_ok = sdot_enabled() && blocked_enabled();
6864        #[cfg(not(any(target_arch = "x86_64", target_arch = "aarch64")))]
6865        let blocked_ok = false;
6866        let run = move |start: usize, end: usize| {
6867            for r in start..end {
6868                let mut bi = 0usize;
6869                #[cfg(target_arch = "aarch64")]
6870                if blocked_ok {
6871                    while bi + 4 <= acts.len() {
6872                        let xs = [
6873                            acts[bi].xq.as_slice(),
6874                            acts[bi + 1].xq.as_slice(),
6875                            acts[bi + 2].xq.as_slice(),
6876                            acts[bi + 3].xq.as_slice(),
6877                        ];
6878                        let d = unsafe { dot_q4t_row_1x4_sdot(bytes, r, gpr, xs) };
6879                        for k in 0..4 {
6880                            let act = &acts[bi + k];
6881                            let mut acc = d[k] * act.sx;
6882                            for &(j, xv) in &act.outliers {
6883                                let (w, sc) = q4t_outlier(bytes, r, gpr, j);
6884                                acc += w * sc * xv;
6885                            }
6886                            // SAFETY: disjoint (bi, r) cells per worker.
6887                            unsafe { *out_addr.at((bi + k) * rows + r) = acc };
6888                        }
6889                        bi += 4;
6890                    }
6891                }
6892                #[cfg(target_arch = "x86_64")]
6893                if blocked_ok {
6894                    while bi + 4 <= acts.len() {
6895                        let xs = [
6896                            acts[bi].xq.as_slice(),
6897                            acts[bi + 1].xq.as_slice(),
6898                            acts[bi + 2].xq.as_slice(),
6899                            acts[bi + 3].xq.as_slice(),
6900                        ];
6901                        let d = unsafe {
6902                            if vnni_tiles_enabled() {
6903                                dot_q4t_row_1x4_vnni(bytes, r, gpr, xs)
6904                            } else {
6905                                dot_q4t_row_1x4_avx2(bytes, r, gpr, xs)
6906                            }
6907                        };
6908                        for k in 0..4 {
6909                            let act = &acts[bi + k];
6910                            let mut acc = d[k] * act.sx;
6911                            for &(j, xv) in &act.outliers {
6912                                let (w, sc) = q4t_outlier(bytes, r, gpr, j);
6913                                acc += w * sc * xv;
6914                            }
6915                            // SAFETY: disjoint (bi, r) cells per worker.
6916                            unsafe { *out_addr.at((bi + k) * rows + r) = acc };
6917                        }
6918                        bi += 4;
6919                    }
6920                }
6921                let _ = blocked_ok;
6922                while bi < acts.len() {
6923                    let act = &acts[bi];
6924                    let mut acc = dot_q4t_row_i8(bytes, r, gpr, &act.xq) * act.sx;
6925                    for &(j, xv) in &act.outliers {
6926                        let (w, s) = q4t_outlier(bytes, r, gpr, j);
6927                        acc += w * s * xv;
6928                    }
6929                    // SAFETY: disjoint (bi, r) cells per worker range.
6930                    unsafe { *out_addr.at(bi * rows + r) = acc };
6931                    bi += 1;
6932                }
6933            }
6934        };
6935        dispatch_rows(pool, rows, &run);
6936        return;
6937    }
6938    let run = move |start: usize, end: usize| {
6939        for r in start..end {
6940            for bi in 0..b {
6941                let x = &xs_all[bi * cols..(bi + 1) * cols];
6942                // SAFETY: disjoint (bi, r) cells per worker range.
6943                unsafe { *out_addr.at(bi * rows + r) = q4t_row_exact(bytes, r, gpr, x) };
6944            }
6945        }
6946    };
6947    dispatch_rows(pool, rows, &run);
6948}
6949
6950// ── q1 (dtype 12): binary weights, [f16 scale][4B sign bits] per
6951// 32-group tile. The kernel family mirrors q4_tiled: one sequential
6952// stream of 6-byte tiles, per-tile integer dot × scale, exact outlier
6953// correction (A8W8 contract), exact scalar path under CMF_SDOT=0. ──
6954
6955/// Per-32-group sums of the quantized activation — the ±1 identity's
6956/// shared half: `dot = −2·sdot(mask, x) − gsum[g]`, computed ONCE per
6957/// matvec and reused by every row.
6958fn q1_group_sums(xq: &[i8], gpr: usize) -> Vec<i32> {
6959    (0..gpr)
6960        .map(|gi| {
6961            xq[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE]
6962                .iter()
6963                .map(|&v| v as i32)
6964                .sum()
6965        })
6966        .collect()
6967}
6968
6969/// One q1 row via the A8W8 int8 path — mask-SDOT on ARM (no ±1
6970/// expansion at all), scalar bit loop elsewhere (AVX2 queued with the
6971/// x86 pass).
6972#[inline]
6973#[allow(unreachable_code)]
6974/// AVX2 q1 row via the same ±1 identity as the ARM sdot kernel: the
6975/// sign bits expand to a {0, −1} byte mask through shuffle+cmpeq, the
6976/// masked activation sums through maddubs(1, x&mask), and
6977/// `dot = −(2·masked_sum + Σx_group)` — bit-identical integer math.
6978#[cfg(target_arch = "x86_64")]
6979#[target_feature(enable = "avx2")]
6980unsafe fn dot_q1_row_avx2(bytes: &[u8], r: usize, gpr: usize, xq: &[i8], gsum: &[i32]) -> f32 {
6981    // SAFETY: callers uphold the 6B-tile and xq/gsum length contracts.
6982    unsafe {
6983        use core::arch::x86_64::*;
6984        // Byte j of the mask must replicate bits-byte j/8.
6985        let expand = _mm256_setr_epi8(
6986            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,
6987            3, 3, 3,
6988        );
6989        let bitsel = _mm256_setr_epi8(
6990            1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64,
6991            -128, 1, 2, 4, 8, 16, 32, 64, -128,
6992        );
6993        let ones8 = _mm256_set1_epi8(1);
6994        let ones16 = _mm256_set1_epi16(1);
6995        let mut acc = 0f32;
6996        for gi in 0..gpr {
6997            let t = bytes.as_ptr().add((r * gpr + gi) * Q1_TILE);
6998            let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
6999            let bits = u32::from_le_bytes([*t.add(2), *t.add(3), *t.add(4), *t.add(5)]);
7000            let bc = _mm256_shuffle_epi8(_mm256_set1_epi32(bits as i32), expand);
7001            let mask = _mm256_cmpeq_epi8(_mm256_and_si256(bc, bitsel), bitsel);
7002            let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
7003            let sel = _mm256_and_si256(x, mask);
7004            // Σ of selected i8 lanes: maddubs(1u8, sel_i8) pairs → madd.
7005            let p16 = _mm256_maddubs_epi16(ones8, sel);
7006            let d32 = _mm256_madd_epi16(p16, ones16);
7007            let hi128 = _mm256_extracti128_si256::<1>(d32);
7008            let s128 = _mm_add_epi32(_mm256_castsi256_si128(d32), hi128);
7009            let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
7010            let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
7011            let msum = _mm_cvtsi128_si32(s32);
7012            // The and-select keeps x UN-negated (unlike ARM's −1-mask
7013            // sdot): d = Σ_set − Σ_unset = 2·Σ_set − Σ_all.
7014            let d = 2 * msum - gsum[gi];
7015            acc += d as f32 * s;
7016        }
7017        acc
7018    }
7019}
7020
7021/// VNNI twin of `dot_q1_row_avx2`: the masked-select sum goes through
7022/// one `vpdpbusd(1u8, sel)` (see `dpbusd_hsum` — bit-identical).
7023#[cfg(target_arch = "x86_64")]
7024#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
7025unsafe fn dot_q1_row_vnni(bytes: &[u8], r: usize, gpr: usize, xq: &[i8], gsum: &[i32]) -> f32 {
7026    // SAFETY: callers uphold the 6B-tile and xq/gsum length contracts.
7027    unsafe {
7028        use core::arch::x86_64::*;
7029        let expand = _mm256_setr_epi8(
7030            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,
7031            3, 3, 3,
7032        );
7033        let bitsel = _mm256_setr_epi8(
7034            1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64,
7035            -128, 1, 2, 4, 8, 16, 32, 64, -128,
7036        );
7037        let ones8 = _mm256_set1_epi8(1);
7038        let mut acc = 0f32;
7039        for gi in 0..gpr {
7040            let t = bytes.as_ptr().add((r * gpr + gi) * Q1_TILE);
7041            let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
7042            let bits = u32::from_le_bytes([*t.add(2), *t.add(3), *t.add(4), *t.add(5)]);
7043            let bc = _mm256_shuffle_epi8(_mm256_set1_epi32(bits as i32), expand);
7044            let mask = _mm256_cmpeq_epi8(_mm256_and_si256(bc, bitsel), bitsel);
7045            let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
7046            let msum = dpbusd_hsum(ones8, _mm256_and_si256(x, mask));
7047            let d = 2 * msum - gsum[gi];
7048            acc += d as f32 * s;
7049        }
7050        acc
7051    }
7052}
7053
7054/// VNNI twin of `dot_q1_row_1x4_avx2` (see `dpbusd_hsum`).
7055#[cfg(target_arch = "x86_64")]
7056#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
7057unsafe fn dot_q1_row_1x4_vnni(
7058    bytes: &[u8],
7059    r: usize,
7060    gpr: usize,
7061    xs: [&[i8]; 4],
7062    gsums: [&[i32]; 4],
7063) -> [f32; 4] {
7064    // SAFETY: callers uphold the 6B-tile and xq/gsum length contracts.
7065    unsafe {
7066        use core::arch::x86_64::*;
7067        let expand = _mm256_setr_epi8(
7068            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,
7069            3, 3, 3,
7070        );
7071        let bitsel = _mm256_setr_epi8(
7072            1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64,
7073            -128, 1, 2, 4, 8, 16, 32, 64, -128,
7074        );
7075        let ones8 = _mm256_set1_epi8(1);
7076        let mut acc = [0f32; 4];
7077        for gi in 0..gpr {
7078            let t = bytes.as_ptr().add((r * gpr + gi) * Q1_TILE);
7079            let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
7080            let bits = u32::from_le_bytes([*t.add(2), *t.add(3), *t.add(4), *t.add(5)]);
7081            let bc = _mm256_shuffle_epi8(_mm256_set1_epi32(bits as i32), expand);
7082            let mask = _mm256_cmpeq_epi8(_mm256_and_si256(bc, bitsel), bitsel);
7083            for (k, xq) in xs.iter().enumerate() {
7084                let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
7085                let msum = dpbusd_hsum(ones8, _mm256_and_si256(x, mask));
7086                let d = 2 * msum - gsums[k][gi];
7087                acc[k] += d as f32 * s;
7088            }
7089        }
7090        acc
7091    }
7092}
7093
7094/// The blocked 1×4 flavor: the expanded bit mask serves four activation
7095/// streams per group (mask build once, four select+reduce chains).
7096#[cfg(target_arch = "x86_64")]
7097#[target_feature(enable = "avx2")]
7098unsafe fn dot_q1_row_1x4_avx2(
7099    bytes: &[u8],
7100    r: usize,
7101    gpr: usize,
7102    xs: [&[i8]; 4],
7103    gsums: [&[i32]; 4],
7104) -> [f32; 4] {
7105    // SAFETY: callers uphold the 6B-tile and xq/gsum length contracts.
7106    unsafe {
7107        use core::arch::x86_64::*;
7108        let expand = _mm256_setr_epi8(
7109            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,
7110            3, 3, 3,
7111        );
7112        let bitsel = _mm256_setr_epi8(
7113            1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64,
7114            -128, 1, 2, 4, 8, 16, 32, 64, -128,
7115        );
7116        let ones8 = _mm256_set1_epi8(1);
7117        let ones16 = _mm256_set1_epi16(1);
7118        let mut acc = [0f32; 4];
7119        for gi in 0..gpr {
7120            let t = bytes.as_ptr().add((r * gpr + gi) * Q1_TILE);
7121            let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
7122            let bits = u32::from_le_bytes([*t.add(2), *t.add(3), *t.add(4), *t.add(5)]);
7123            let bc = _mm256_shuffle_epi8(_mm256_set1_epi32(bits as i32), expand);
7124            let mask = _mm256_cmpeq_epi8(_mm256_and_si256(bc, bitsel), bitsel);
7125            for (k, xq) in xs.iter().enumerate() {
7126                let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
7127                let sel = _mm256_and_si256(x, mask);
7128                let p16 = _mm256_maddubs_epi16(ones8, sel);
7129                let d32 = _mm256_madd_epi16(p16, ones16);
7130                let hi128 = _mm256_extracti128_si256::<1>(d32);
7131                let s128 = _mm_add_epi32(_mm256_castsi256_si128(d32), hi128);
7132                let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
7133                let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
7134                let msum = _mm_cvtsi128_si32(s32);
7135                let d = 2 * msum - gsums[k][gi];
7136                acc[k] += d as f32 * s;
7137            }
7138        }
7139        acc
7140    }
7141}
7142
7143#[allow(unreachable_code)]
7144fn dot_q1_row_i8(bytes: &[u8], r: usize, gpr: usize, xq: &[i8], gsum: &[i32]) -> f32 {
7145    #[cfg(target_arch = "aarch64")]
7146    unsafe {
7147        return dot_q1_row_sdot(bytes, r, gpr, xq, gsum);
7148    }
7149    #[cfg(target_arch = "x86_64")]
7150    if avx2_enabled() {
7151        unsafe {
7152            if vnni_tiles_enabled() {
7153                return dot_q1_row_vnni(bytes, r, gpr, xq, gsum);
7154            }
7155            return dot_q1_row_avx2(bytes, r, gpr, xq, gsum);
7156        }
7157    }
7158    let _ = gsum;
7159    let mut acc = 0f32;
7160    for gi in 0..gpr {
7161        let tile = &bytes[(r * gpr + gi) * Q1_TILE..(r * gpr + gi + 1) * Q1_TILE];
7162        let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
7163        let mut d = 0i32;
7164        for (j, &b) in tile[2..].iter().enumerate() {
7165            for k in 0..8 {
7166                let w = ((b >> k) & 1) as i32 * 2 - 1;
7167                d += w * xq[gi * GROUP_SIZE + j * 8 + k] as i32;
7168            }
7169        }
7170        acc += d as f32 * s;
7171    }
7172    acc
7173}
7174
7175/// SDOT q1 row via the ±1 identity: the vtst mask (0xFF where the bit
7176/// is set, i.e. −1 as i8) feeds `sdot` DIRECTLY — no expansion to ±1
7177/// lanes at all — and `dot = −(2·sdot(mask, x) + Σx_group)`, with the
7178/// per-group activation sums shared across every row of the matvec.
7179/// Four tiles (128 weights) per iteration: integer dots reduce through
7180/// a vpaddq tree into ONE i32x4 that meets its four scales in a single
7181/// fused f32 multiply-add. Integer math throughout — bit-identical to
7182/// the scalar ±1 reference.
7183#[cfg(target_arch = "aarch64")]
7184#[target_feature(enable = "neon,dotprod")]
7185unsafe fn dot_q1_row_sdot(bytes: &[u8], r: usize, gpr: usize, xq: &[i8], gsum: &[i32]) -> f32 {
7186    // SAFETY: callers uphold slice-length contracts (6B tile per group,
7187    // xq.len() == gpr·GROUP_SIZE, gsum.len() == gpr).
7188    unsafe {
7189        use core::arch::aarch64::*;
7190        use core::arch::asm;
7191        const MASKS: [u8; 16] = [1, 2, 4, 8, 16, 32, 64, 128, 1, 2, 4, 8, 16, 32, 64, 128];
7192        let m = vld1q_u8(MASKS.as_ptr());
7193        // One tile's −Σ_set(x) as an UNREDUCED i32x4 (two mask-sdots).
7194        macro_rules! tile_dot {
7195            ($t:expr, $x:expr) => {{
7196                let v0 = vcombine_u8(vdup_n_u8(*$t.add(2)), vdup_n_u8(*$t.add(3)));
7197                let v1 = vcombine_u8(vdup_n_u8(*$t.add(4)), vdup_n_u8(*$t.add(5)));
7198                let w0 = vreinterpretq_s8_u8(vtstq_u8(v0, m));
7199                let w1 = vreinterpretq_s8_u8(vtstq_u8(v1, m));
7200                let x0 = vld1q_s8($x);
7201                let x1 = vld1q_s8($x.add(16));
7202                let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
7203                asm!(
7204                    "sdot {a0:v}.4s, {w0:v}.16b, {x0:v}.16b",
7205                    "sdot {a1:v}.4s, {w1:v}.16b, {x1:v}.16b",
7206                    a0 = inout(vreg) a0, a1 = inout(vreg) a1,
7207                    w0 = in(vreg) w0, x0 = in(vreg) x0, w1 = in(vreg) w1, x1 = in(vreg) x1,
7208                    options(pure, nomem, nostack),
7209                );
7210                vaddq_s32(a0, a1)
7211            }};
7212        }
7213        // TBL unpack over PAIR loads: one vld1q covers two 6B tiles
7214        // ([s s b b b b][s s b b b b] + 4B slack), TBL replicates each
7215        // bit-byte across 8 lanes for vtst, and the four scales gather
7216        // through tbl2 into one fcvtl — the 16 ld1r broadcast loads and
7217        // 4 branchy software f16 conversions per 128 weights (the
7218        // measured load-port wall of this kernel) become 2 vector
7219        // loads + 9 table lookups. Integer math order is unchanged —
7220        // bit-identical results (FCVTL is exact on every f16).
7221        const IW00: [u8; 16] = [2, 2, 2, 2, 2, 2, 2, 2, 3, 3, 3, 3, 3, 3, 3, 3];
7222        const IW01: [u8; 16] = [4, 4, 4, 4, 4, 4, 4, 4, 5, 5, 5, 5, 5, 5, 5, 5];
7223        const IW10: [u8; 16] = [8, 8, 8, 8, 8, 8, 8, 8, 9, 9, 9, 9, 9, 9, 9, 9];
7224        const IW11: [u8; 16] = [
7225            10, 10, 10, 10, 10, 10, 10, 10, 11, 11, 11, 11, 11, 11, 11, 11,
7226        ];
7227        const ISC: [u8; 8] = [0, 1, 6, 7, 16, 17, 22, 23];
7228        let (iw00, iw01) = (vld1q_u8(IW00.as_ptr()), vld1q_u8(IW01.as_ptr()));
7229        let (iw10, iw11) = (vld1q_u8(IW10.as_ptr()), vld1q_u8(IW11.as_ptr()));
7230        let isc = vld1_u8(ISC.as_ptr());
7231        // One tile's −Σ_set(x) from a TBL-unpacked pair load.
7232        macro_rules! tile_dot_tbl {
7233            ($ld:expr, $i0:expr, $i1:expr, $x:expr) => {{
7234                let w0 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8($ld, $i0), m));
7235                let w1 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8($ld, $i1), m));
7236                let x0 = vld1q_s8($x);
7237                let x1 = vld1q_s8($x.add(16));
7238                let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
7239                asm!(
7240                    "sdot {a0:v}.4s, {w0:v}.16b, {x0:v}.16b",
7241                    "sdot {a1:v}.4s, {w1:v}.16b, {x1:v}.16b",
7242                    a0 = inout(vreg) a0, a1 = inout(vreg) a1,
7243                    w0 = in(vreg) w0, x0 = in(vreg) x0, w1 = in(vreg) w1, x1 = in(vreg) x1,
7244                    options(pure, nomem, nostack),
7245                );
7246                vaddq_s32(a0, a1)
7247            }};
7248        }
7249        let base = bytes.as_ptr().add(r * gpr * Q1_TILE);
7250        let row_base = r * gpr * Q1_TILE;
7251        let abs_end = bytes.len();
7252        let xp = xq.as_ptr();
7253        let gp = gsum.as_ptr();
7254        let mut accv = vdupq_n_f32(0.0);
7255        let mut gi = 0;
7256        // The second pair load reads 4B past tile gi+3 — stay inside
7257        // the payload slice (only the file's final tiles fall back).
7258        while gi + 4 <= gpr && row_base + (gi + 4) * Q1_TILE + 4 <= abs_end {
7259            let t0 = base.add(gi * Q1_TILE);
7260            let ld_a = vld1q_u8(t0);
7261            let ld_b = vld1q_u8(t0.add(2 * Q1_TILE));
7262            let d0 = tile_dot_tbl!(ld_a, iw00, iw01, xp.add(gi * GROUP_SIZE));
7263            let d1 = tile_dot_tbl!(ld_a, iw10, iw11, xp.add((gi + 1) * GROUP_SIZE));
7264            let d2 = tile_dot_tbl!(ld_b, iw00, iw01, xp.add((gi + 2) * GROUP_SIZE));
7265            let d3 = tile_dot_tbl!(ld_b, iw10, iw11, xp.add((gi + 3) * GROUP_SIZE));
7266            // [−Σ0, −Σ1, −Σ2, −Σ3] → dots = −(2·Σset_neg + gsum)
7267            let neg = vpaddq_s32(vpaddq_s32(d0, d1), vpaddq_s32(d2, d3));
7268            let g = vld1q_s32(gp.add(gi));
7269            let dots = vnegq_s32(vaddq_s32(vshlq_n_s32::<1>(neg), g));
7270            let sc16 = vqtbl2_u8(uint8x16x2_t(ld_a, ld_b), isc);
7271            let scf: float32x4_t;
7272            asm!(
7273                "fcvtl {o:v}.4s, {i:v}.4h",
7274                o = out(vreg) scf, i = in(vreg) sc16,
7275                options(pure, nomem, nostack),
7276            );
7277            accv = vfmaq_f32(accv, vcvtq_f32_s32(dots), scf);
7278            gi += 4;
7279        }
7280        let mut acc = vaddvq_f32(accv);
7281        while gi < gpr {
7282            let t = base.add(gi * Q1_TILE);
7283            let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
7284            let d = vaddvq_s32(tile_dot!(t, xp.add(gi * GROUP_SIZE)));
7285            acc += (-(2 * d + *gp.add(gi))) as f32 * s;
7286            gi += 1;
7287        }
7288        acc
7289    }
7290}
7291
7292/// Blocked q1 1×4: one TBL unpack of the tile pair serves FOUR
7293/// activation streams (prefill amortization — the same idea as the
7294/// AVX2 twin; per stream the group order, fma order and tail match the
7295/// single-row kernel exactly, so batch == matvec bit-for-bit).
7296#[cfg(target_arch = "aarch64")]
7297#[target_feature(enable = "neon,dotprod")]
7298unsafe fn dot_q1_row_1x4_sdot(
7299    bytes: &[u8],
7300    r: usize,
7301    gpr: usize,
7302    xs: [&[i8]; 4],
7303    gs: [&[i32]; 4],
7304) -> [f32; 4] {
7305    // SAFETY: same slice-length contracts as `dot_q1_row_sdot`, ×4.
7306    unsafe {
7307        use core::arch::aarch64::*;
7308        use core::arch::asm;
7309        const MASKS: [u8; 16] = [1, 2, 4, 8, 16, 32, 64, 128, 1, 2, 4, 8, 16, 32, 64, 128];
7310        const IW00: [u8; 16] = [2, 2, 2, 2, 2, 2, 2, 2, 3, 3, 3, 3, 3, 3, 3, 3];
7311        const IW01: [u8; 16] = [4, 4, 4, 4, 4, 4, 4, 4, 5, 5, 5, 5, 5, 5, 5, 5];
7312        const IW10: [u8; 16] = [8, 8, 8, 8, 8, 8, 8, 8, 9, 9, 9, 9, 9, 9, 9, 9];
7313        const IW11: [u8; 16] = [
7314            10, 10, 10, 10, 10, 10, 10, 10, 11, 11, 11, 11, 11, 11, 11, 11,
7315        ];
7316        const ISC: [u8; 8] = [0, 1, 6, 7, 16, 17, 22, 23];
7317        let m = vld1q_u8(MASKS.as_ptr());
7318        let (iw00, iw01) = (vld1q_u8(IW00.as_ptr()), vld1q_u8(IW01.as_ptr()));
7319        let (iw10, iw11) = (vld1q_u8(IW10.as_ptr()), vld1q_u8(IW11.as_ptr()));
7320        let isc = vld1_u8(ISC.as_ptr());
7321        macro_rules! sdot2 {
7322            ($w0:expr, $w1:expr, $x:expr) => {{
7323                let x0 = vld1q_s8($x);
7324                let x1 = vld1q_s8($x.add(16));
7325                let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
7326                asm!(
7327                    "sdot {a0:v}.4s, {w0:v}.16b, {x0:v}.16b",
7328                    "sdot {a1:v}.4s, {w1:v}.16b, {x1:v}.16b",
7329                    a0 = inout(vreg) a0, a1 = inout(vreg) a1,
7330                    w0 = in(vreg) $w0, x0 = in(vreg) x0, w1 = in(vreg) $w1, x1 = in(vreg) x1,
7331                    options(pure, nomem, nostack),
7332                );
7333                vaddq_s32(a0, a1)
7334            }};
7335        }
7336        let base = bytes.as_ptr().add(r * gpr * Q1_TILE);
7337        let row_base = r * gpr * Q1_TILE;
7338        let abs_end = bytes.len();
7339        let mut accv = [vdupq_n_f32(0.0); 4];
7340        let mut gi = 0;
7341        while gi + 4 <= gpr && row_base + (gi + 4) * Q1_TILE + 4 <= abs_end {
7342            let t0 = base.add(gi * Q1_TILE);
7343            let ld_a = vld1q_u8(t0);
7344            let ld_b = vld1q_u8(t0.add(2 * Q1_TILE));
7345            // Unpack ONCE — eight ±mask vectors serve all four streams.
7346            let w00 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8(ld_a, iw00), m));
7347            let w01 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8(ld_a, iw01), m));
7348            let w10 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8(ld_a, iw10), m));
7349            let w11 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8(ld_a, iw11), m));
7350            let w20 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8(ld_b, iw00), m));
7351            let w21 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8(ld_b, iw01), m));
7352            let w30 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8(ld_b, iw10), m));
7353            let w31 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8(ld_b, iw11), m));
7354            let sc16 = vqtbl2_u8(uint8x16x2_t(ld_a, ld_b), isc);
7355            let scf: float32x4_t;
7356            asm!(
7357                "fcvtl {o:v}.4s, {i:v}.4h",
7358                o = out(vreg) scf, i = in(vreg) sc16,
7359                options(pure, nomem, nostack),
7360            );
7361            for k in 0..4 {
7362                let xp = xs[k].as_ptr();
7363                let d0 = sdot2!(w00, w01, xp.add(gi * GROUP_SIZE));
7364                let d1 = sdot2!(w10, w11, xp.add((gi + 1) * GROUP_SIZE));
7365                let d2 = sdot2!(w20, w21, xp.add((gi + 2) * GROUP_SIZE));
7366                let d3 = sdot2!(w30, w31, xp.add((gi + 3) * GROUP_SIZE));
7367                let neg = vpaddq_s32(vpaddq_s32(d0, d1), vpaddq_s32(d2, d3));
7368                let g = vld1q_s32(gs[k].as_ptr().add(gi));
7369                let dots = vnegq_s32(vaddq_s32(vshlq_n_s32::<1>(neg), g));
7370                accv[k] = vfmaq_f32(accv[k], vcvtq_f32_s32(dots), scf);
7371            }
7372            gi += 4;
7373        }
7374        let mut acc = [
7375            vaddvq_f32(accv[0]),
7376            vaddvq_f32(accv[1]),
7377            vaddvq_f32(accv[2]),
7378            vaddvq_f32(accv[3]),
7379        ];
7380        while gi < gpr {
7381            let t = base.add(gi * Q1_TILE);
7382            let sc = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
7383            let v0 = vcombine_u8(vdup_n_u8(*t.add(2)), vdup_n_u8(*t.add(3)));
7384            let v1 = vcombine_u8(vdup_n_u8(*t.add(4)), vdup_n_u8(*t.add(5)));
7385            let w0 = vreinterpretq_s8_u8(vtstq_u8(v0, m));
7386            let w1 = vreinterpretq_s8_u8(vtstq_u8(v1, m));
7387            for k in 0..4 {
7388                let d = vaddvq_s32(sdot2!(w0, w1, xs[k].as_ptr().add(gi * GROUP_SIZE)));
7389                acc[k] += (-(2 * d + *gs[k].as_ptr().add(gi))) as f32 * sc;
7390            }
7391            gi += 1;
7392        }
7393        acc
7394    }
7395}
7396
7397/// (weight ±1, scale) of one q1 element — the exact outlier term.
7398#[inline]
7399fn q1_outlier(bytes: &[u8], r: usize, gpr: usize, j: usize) -> (f32, f32) {
7400    let gi = j / GROUP_SIZE;
7401    let k = j % GROUP_SIZE;
7402    let tile = &bytes[(r * gpr + gi) * Q1_TILE..(r * gpr + gi + 1) * Q1_TILE];
7403    let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
7404    let bit = (tile[2 + k / 8] >> (k % 8)) & 1;
7405    ((bit as i32 * 2 - 1) as f32, s)
7406}
7407
7408/// Exact scalar q1 row (CMF_SDOT=0 contract).
7409#[inline]
7410fn q1_row_exact(bytes: &[u8], r: usize, gpr: usize, x: &[f32]) -> f32 {
7411    let mut acc = 0f32;
7412    for gi in 0..gpr {
7413        let tile = &bytes[(r * gpr + gi) * Q1_TILE..(r * gpr + gi + 1) * Q1_TILE];
7414        let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
7415        let xg = &x[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
7416        let mut ga = 0f32;
7417        for (j, &b) in tile[2..].iter().enumerate() {
7418            for k in 0..8 {
7419                ga += (((b >> k) & 1) as f32 * 2.0 - 1.0) * xg[j * 8 + k];
7420            }
7421        }
7422        acc += ga * s;
7423    }
7424    acc
7425}
7426
7427/// One q1 row range via A8W8 (the body of `q1_matvec`'s hot loop,
7428/// extracted so multi-matrix jobs drive the same kernel).
7429#[allow(clippy::too_many_arguments)]
7430fn q1_range_a8w8(
7431    bytes: &[u8],
7432    gpr: usize,
7433    act: &SplitAct,
7434    gsum: &[i32],
7435    out: SendMut,
7436    start: usize,
7437    end: usize,
7438) {
7439    for r in start..end {
7440        let mut acc = dot_q1_row_i8(bytes, r, gpr, &act.xq, gsum) * act.sx;
7441        for &(j, xv) in &act.outliers {
7442            let (w, s) = q1_outlier(bytes, r, gpr, j);
7443            acc += w * s * xv;
7444        }
7445        // SAFETY: disjoint row ranges per worker.
7446        unsafe { *out.at(r) = acc };
7447    }
7448}
7449
7450/// Exact-scalar q1 row range (CMF_SDOT=0 contract).
7451fn q1_range_f32(bytes: &[u8], gpr: usize, x: &[f32], out: SendMut, start: usize, end: usize) {
7452    for r in start..end {
7453        // SAFETY: disjoint row ranges per worker.
7454        unsafe { *out.at(r) = q1_row_exact(bytes, r, gpr, x) };
7455    }
7456}
7457
7458/// q1t per-row overlay locator. After the base (`base_len`) come
7459/// `[u32 row_ptr[rows+1]]` then `[(u16 col, f16 val)]` grouped by row (row
7460/// `r`'s entries are `[row_ptr[r], row_ptr[r+1])`). Returns
7461/// `(row_ptr offset, entries offset, present)`.
7462fn q1t_overlay(bytes: &[u8], base_len: usize, rows: usize) -> (usize, usize, bool) {
7463    let entries = base_len + (rows + 1) * 4;
7464    (base_len, entries, entries <= bytes.len())
7465}
7466
7467/// Read `row_ptr[r]` from the overlay's prefix-sum table.
7468#[inline]
7469fn q1t_rowptr(bytes: &[u8], rp_off: usize, r: usize) -> usize {
7470    let o = rp_off + r * 4;
7471    u32::from_le_bytes([bytes[o], bytes[o + 1], bytes[o + 2], bytes[o + 3]]) as usize
7472}
7473
7474/// Byte → the 5 ternary signs it packs `{−1,0,+1}` as f32, precomputed so
7475/// decoding a q1t code is a table load, not the base-3 divide/modulo per
7476/// weight (division is ~20–40× the cost of a load). Built at compile time.
7477const SIGN5: [[f32; 5]; 256] = {
7478    let mut lut = [[0.0f32; 5]; 256];
7479    let pow3 = [1u16, 3, 9, 27, 81];
7480    let mut byte = 0usize;
7481    while byte < 256 {
7482        let mut i = 0usize;
7483        while i < 5 {
7484            let code = (byte as u16 / pow3[i]) % 3;
7485            lut[byte][i] = if code == 1 {
7486                1.0
7487            } else if code == 2 {
7488                -1.0
7489            } else {
7490                0.0
7491            };
7492            i += 1;
7493        }
7494        byte += 1;
7495    }
7496    lut
7497};
7498
7499/// Same table, as i8 signs — the operand for the int8 SDOT base kernel.
7500const SIGN5_I8: [[i8; 5]; 256] = {
7501    let mut lut = [[0i8; 5]; 256];
7502    let pow3 = [1u16, 3, 9, 27, 81];
7503    let mut byte = 0usize;
7504    while byte < 256 {
7505        let mut i = 0usize;
7506        while i < 5 {
7507            let code = (byte as u16 / pow3[i]) % 3;
7508            lut[byte][i] = if code == 1 {
7509                1
7510            } else if code == 2 {
7511                -1
7512            } else {
7513                0
7514            };
7515            i += 1;
7516        }
7517        byte += 1;
7518    }
7519    lut
7520};
7521
7522/// The same 5 i8 signs packed into a u64 (`[s0 s1 s2 s3 s4 0 0 0]`, LE) so the
7523/// group unpack is 7 unaligned u64 stores at offsets 0,5,10,…,30 instead of
7524/// six 5-byte copies + LUT indexing — each store's trailing zeros are fixed by
7525/// the next store, and the last one runs 6 B past the 32nd weight (the unpack
7526/// buffer is padded to 40). This is the decode/prefill hot inner op.
7527const SIGN5_U64: [u64; 256] = {
7528    let mut lut = [0u64; 256];
7529    let pow3 = [1u16, 3, 9, 27, 81];
7530    let mut byte = 0usize;
7531    while byte < 256 {
7532        let mut v = 0u64;
7533        let mut i = 0usize;
7534        while i < 5 {
7535            let code = (byte as u16 / pow3[i]) % 3;
7536            let s: u8 = if code == 1 {
7537                1
7538            } else if code == 2 {
7539                0xFF
7540            } else {
7541                0
7542            };
7543            v |= (s as u64) << (i * 8);
7544            i += 1;
7545        }
7546        lut[byte] = v;
7547        byte += 1;
7548    }
7549    lut
7550};
7551
7552/// Ternary base weight at `(row r, col j)` = `sign(code)·s_group`. Used to add
7553/// back activation-outlier columns, whose `x` was zeroed for the int8 bulk dot
7554/// (`split_act`). At a weight-outlier position the code is 0, so this is 0 and
7555/// the overlay correction owns that column — no double counting.
7556#[inline]
7557fn q1t_base_weight(bytes: &[u8], r: usize, gpr: usize, j: usize) -> f32 {
7558    const TILE: usize = cortiq_core::quant::Q1T_TILE;
7559    let off = (r * gpr + j / GROUP_SIZE) * TILE;
7560    let s = f16_to_f32(u16::from_le_bytes([bytes[off], bytes[off + 1]]));
7561    let within = j % GROUP_SIZE;
7562    SIGN5[bytes[off + 2 + within / 5] as usize][within % 5] * s
7563}
7564
7565/// One 32-group int8 dot via two SDOTs. Bit-exact vs the scalar i8 sum
7566/// (integer accumulation is order-independent).
7567#[cfg(target_arch = "aarch64")]
7568#[target_feature(enable = "neon,dotprod")]
7569#[inline]
7570unsafe fn sdot32_i8(w: *const i8, x: *const i8) -> i32 {
7571    // SAFETY: caller guarantees 32 readable i8 at each pointer.
7572    unsafe {
7573        use core::arch::aarch64::*;
7574        use core::arch::asm;
7575        let w0 = vld1q_s8(w);
7576        let w1 = vld1q_s8(w.add(16));
7577        let x0 = vld1q_s8(x);
7578        let x1 = vld1q_s8(x.add(16));
7579        let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
7580        asm!(
7581            "sdot {a0:v}.4s, {w0:v}.16b, {x0:v}.16b",
7582            "sdot {a1:v}.4s, {w1:v}.16b, {x1:v}.16b",
7583            a0 = inout(vreg) a0, a1 = inout(vreg) a1,
7584            w0 = in(vreg) w0, x0 = in(vreg) x0, w1 = in(vreg) w1, x1 = in(vreg) x1,
7585            options(pure, nomem, nostack),
7586        );
7587        vaddvq_s32(vaddq_s32(a0, a1))
7588    }
7589}
7590
7591/// One 32-group int8 dot via AVX2: signed·signed as `maddubs(|w|, sign(x,w))`
7592/// then `madd` and a horizontal reduce (the same idiom as `dot_q4t_row_avx2`).
7593#[cfg(target_arch = "x86_64")]
7594#[target_feature(enable = "avx2")]
7595#[inline]
7596unsafe fn i8dot32_avx2(w: *const i8, x: *const i8) -> i32 {
7597    // SAFETY: caller guarantees 32 readable i8 at each pointer.
7598    unsafe {
7599        use core::arch::x86_64::*;
7600        let wv = _mm256_loadu_si256(w as *const __m256i);
7601        let xv = _mm256_loadu_si256(x as *const __m256i);
7602        let p16 = _mm256_maddubs_epi16(_mm256_abs_epi8(wv), _mm256_sign_epi8(xv, wv));
7603        let d = _mm256_madd_epi16(p16, _mm256_set1_epi16(1));
7604        let hi128 = _mm256_extracti128_si256::<1>(d);
7605        let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
7606        let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
7607        let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
7608        _mm_cvtsi128_si32(s32)
7609    }
7610}
7611
7612/// Unpack one q1t group's base-3 codes into 32 i8 signs via 7 unaligned u64
7613/// stores (see `SIGN5_U64`). `dst` MUST have ≥ 40 bytes: the 7th store writes
7614/// `dst[30..38]`. Stores go in order so each one's trailing zeros are
7615/// overwritten by the next; the final 6 padding bytes are unused by the dot.
7616#[inline]
7617fn q1t_unpack_group_i8(codes: *const u8, dst: &mut [i8]) {
7618    debug_assert!(dst.len() >= 40);
7619    // SAFETY: codes points at 7 readable bytes; dst has ≥ 40 bytes so every
7620    // 8-byte store at offset bi*5 (bi ≤ 6 → ≤ 30) stays in bounds.
7621    unsafe {
7622        let p = dst.as_mut_ptr();
7623        for bi in 0..7 {
7624            core::ptr::write_unaligned(
7625                p.add(bi * 5) as *mut u64,
7626                SIGN5_U64[*codes.add(bi) as usize],
7627            );
7628        }
7629    }
7630}
7631
7632/// One 32-group int8 dot, arch-dispatched (the matmat inner loop, where the
7633/// row's signs are unpacked once and dotted against every batch input).
7634/// Callers are gated by `a8w8_enabled()`, so the target-feature arms are
7635/// reachable; the scalar arm is a non-SIMD-arch fallback.
7636#[inline]
7637fn q1t_i8dot32(w: *const i8, x: *const i8) -> i32 {
7638    #[cfg(target_arch = "aarch64")]
7639    unsafe {
7640        return sdot32_i8(w, x);
7641    }
7642    #[cfg(target_arch = "x86_64")]
7643    unsafe {
7644        return i8dot32_avx2(w, x);
7645    }
7646    #[allow(unreachable_code)]
7647    unsafe {
7648        let mut s = 0i32;
7649        for k in 0..GROUP_SIZE {
7650            s += *w.add(k) as i32 * *x.add(k) as i32;
7651        }
7652        s
7653    }
7654}
7655
7656#[inline]
7657unsafe fn q1t_unpack_reg_u64s(codes: *const u8) -> (u64, u64, u64, u64) {
7658    let (s0, s1, s2, s3, s4, s5, s6) = unsafe {
7659        (
7660            SIGN5_U64[*codes as usize],
7661            SIGN5_U64[*codes.add(1) as usize],
7662            SIGN5_U64[*codes.add(2) as usize],
7663            SIGN5_U64[*codes.add(3) as usize],
7664            SIGN5_U64[*codes.add(4) as usize],
7665            SIGN5_U64[*codes.add(5) as usize],
7666            SIGN5_U64[*codes.add(6) as usize],
7667        )
7668    };
7669
7670    let u0 = s0 | (s1 << 40);
7671    let u1 = (s1 >> 24) | (s2 << 16) | (s3 << 56);
7672    let u2 = (s3 >> 8) | (s4 << 32);
7673    let u3 = (s4 >> 32) | (s5 << 8) | (s6 << 48);
7674
7675    (u0, u1, u2, u3)
7676}
7677
7678/// One q1t row's int8 base dot: `Σ_group s·dot(signs, xq)` (before the shared
7679/// `sx`). Direct register unpacking (zero stack stores/loads, no STLF stalls).
7680/// ARM SDOT.
7681#[cfg(target_arch = "aarch64")]
7682#[target_feature(enable = "neon,dotprod")]
7683unsafe fn q1t_dot_row_sdot(bytes: &[u8], r: usize, gpr: usize, xq: &[i8]) -> f32 {
7684    use core::arch::aarch64::*;
7685    use core::arch::asm;
7686    unsafe {
7687        const TILE: usize = cortiq_core::quant::Q1T_TILE;
7688        let mut acc = 0f32;
7689        let bytes_ptr = bytes.as_ptr();
7690        let xq_ptr = xq.as_ptr();
7691        let row_off = r * gpr * TILE;
7692
7693        let gpr2 = gpr & !1;
7694        let mut gi = 0;
7695        while gi < gpr2 {
7696            let off0 = row_off + gi * TILE;
7697            let off1 = off0 + TILE;
7698            let s0 = f16_to_f32(u16::from_le_bytes([
7699                *bytes_ptr.add(off0),
7700                *bytes_ptr.add(off0 + 1),
7701            ]));
7702            let s1 = f16_to_f32(u16::from_le_bytes([
7703                *bytes_ptr.add(off1),
7704                *bytes_ptr.add(off1 + 1),
7705            ]));
7706
7707            let (u0_0, u1_0, u2_0, u3_0) = q1t_unpack_reg_u64s(bytes_ptr.add(off0 + 2));
7708            let (u0_1, u1_1, u2_1, u3_1) = q1t_unpack_reg_u64s(bytes_ptr.add(off1 + 2));
7709
7710            let w0_0 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u0_0), vcreate_u64(u1_0)));
7711            let w1_0 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u2_0), vcreate_u64(u3_0)));
7712            let w0_1 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u0_1), vcreate_u64(u1_1)));
7713            let w1_1 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u2_1), vcreate_u64(u3_1)));
7714
7715            let x0_0 = vld1q_s8(xq_ptr.add(gi * GROUP_SIZE));
7716            let x1_0 = vld1q_s8(xq_ptr.add(gi * GROUP_SIZE + 16));
7717            let x0_1 = vld1q_s8(xq_ptr.add((gi + 1) * GROUP_SIZE));
7718            let x1_1 = vld1q_s8(xq_ptr.add((gi + 1) * GROUP_SIZE + 16));
7719
7720            let (mut a0_0, mut a1_0) = (vdupq_n_s32(0), vdupq_n_s32(0));
7721            let (mut a0_1, mut a1_1) = (vdupq_n_s32(0), vdupq_n_s32(0));
7722            asm!(
7723                "sdot {a0_0:v}.4s, {w0_0:v}.16b, {x0_0:v}.16b",
7724                "sdot {a1_0:v}.4s, {w1_0:v}.16b, {x1_0:v}.16b",
7725                "sdot {a0_1:v}.4s, {w0_1:v}.16b, {x0_1:v}.16b",
7726                "sdot {a1_1:v}.4s, {w1_1:v}.16b, {x1_1:v}.16b",
7727                a0_0 = inout(vreg) a0_0, a1_0 = inout(vreg) a1_0,
7728                a0_1 = inout(vreg) a0_1, a1_1 = inout(vreg) a1_1,
7729                w0_0 = in(vreg) w0_0, x0_0 = in(vreg) x0_0, w1_0 = in(vreg) w1_0, x1_0 = in(vreg) x1_0,
7730                w0_1 = in(vreg) w0_1, x0_1 = in(vreg) x0_1, w1_1 = in(vreg) w1_1, x1_1 = in(vreg) x1_1,
7731                options(pure, nomem, nostack),
7732            );
7733            let d0 = vaddvq_s32(vaddq_s32(a0_0, a1_0));
7734            let d1 = vaddvq_s32(vaddq_s32(a0_1, a1_1));
7735            acc += d0 as f32 * s0 + d1 as f32 * s1;
7736            gi += 2;
7737        }
7738
7739        if gi < gpr {
7740            let off = row_off + gi * TILE;
7741            let s = f16_to_f32(u16::from_le_bytes([
7742                *bytes_ptr.add(off),
7743                *bytes_ptr.add(off + 1),
7744            ]));
7745            let (u0, u1, u2, u3) = q1t_unpack_reg_u64s(bytes_ptr.add(off + 2));
7746            let w0 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u0), vcreate_u64(u1)));
7747            let w1 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u2), vcreate_u64(u3)));
7748            let x0 = vld1q_s8(xq_ptr.add(gi * GROUP_SIZE));
7749            let x1 = vld1q_s8(xq_ptr.add(gi * GROUP_SIZE + 16));
7750            let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
7751            asm!(
7752                "sdot {a0:v}.4s, {w0:v}.16b, {x0:v}.16b",
7753                "sdot {a1:v}.4s, {w1:v}.16b, {x1:v}.16b",
7754                a0 = inout(vreg) a0, a1 = inout(vreg) a1,
7755                w0 = in(vreg) w0, x0 = in(vreg) x0, w1 = in(vreg) w1, x1 = in(vreg) x1,
7756                options(pure, nomem, nostack),
7757            );
7758            let d = vaddvq_s32(vaddq_s32(a0, a1));
7759            acc += d as f32 * s;
7760        }
7761        acc
7762    }
7763}
7764
7765/// x86 AVX2 mirror of `q1t_dot_row_sdot` (maddubs int8 dot per group).
7766#[cfg(target_arch = "x86_64")]
7767#[target_feature(enable = "avx2")]
7768unsafe fn q1t_dot_row_avx2(bytes: &[u8], r: usize, gpr: usize, xq: &[i8]) -> f32 {
7769    use core::arch::x86_64::*;
7770    unsafe {
7771        const TILE: usize = cortiq_core::quant::Q1T_TILE;
7772        let mut acc = 0f32;
7773        let bytes_ptr = bytes.as_ptr();
7774        let xq_ptr = xq.as_ptr();
7775        let row_off = r * gpr * TILE;
7776
7777        let ones = _mm256_set1_epi16(1);
7778        for gi in 0..gpr {
7779            let off = row_off + gi * TILE;
7780            let s = f16_to_f32(u16::from_le_bytes([
7781                *bytes_ptr.add(off),
7782                *bytes_ptr.add(off + 1),
7783            ]));
7784            let (u0, u1, u2, u3) = q1t_unpack_reg_u64s(bytes_ptr.add(off + 2));
7785            let wv = _mm256_set_epi64x(u3 as i64, u2 as i64, u1 as i64, u0 as i64);
7786            let xv = _mm256_loadu_si256(xq_ptr.add(gi * GROUP_SIZE) as *const __m256i);
7787            let p16 = _mm256_maddubs_epi16(_mm256_abs_epi8(wv), _mm256_sign_epi8(xv, wv));
7788            let d256 = _mm256_madd_epi16(p16, ones);
7789            let d128 = _mm_add_epi32(
7790                _mm256_castsi256_si128(d256),
7791                _mm256_extracti128_si256(d256, 1),
7792            );
7793            let d64 = _mm_add_epi32(d128, _mm_shuffle_epi32(d128, 0xee));
7794            let d32 = _mm_cvtsi128_si32(_mm_add_epi32(d64, _mm_shuffle_epi32(d64, 0x55)));
7795            acc += d32 as f32 * s;
7796        }
7797        acc
7798    }
7799}
7800
7801/// VNNI twin of `q1t_dot_row_avx2` (see `dpbusd_hsum`).
7802#[cfg(target_arch = "x86_64")]
7803#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
7804unsafe fn q1t_dot_row_vnni(bytes: &[u8], r: usize, gpr: usize, xq: &[i8]) -> f32 {
7805    use core::arch::x86_64::*;
7806    // SAFETY: same tile/xq contracts as `q1t_dot_row_avx2`.
7807    unsafe {
7808        const TILE: usize = cortiq_core::quant::Q1T_TILE;
7809        let mut acc = 0f32;
7810        let bytes_ptr = bytes.as_ptr();
7811        let xq_ptr = xq.as_ptr();
7812        let row_off = r * gpr * TILE;
7813        for gi in 0..gpr {
7814            let off = row_off + gi * TILE;
7815            let s = f16_to_f32(u16::from_le_bytes([
7816                *bytes_ptr.add(off),
7817                *bytes_ptr.add(off + 1),
7818            ]));
7819            let (u0, u1, u2, u3) = q1t_unpack_reg_u64s(bytes_ptr.add(off + 2));
7820            let wv = _mm256_set_epi64x(u3 as i64, u2 as i64, u1 as i64, u0 as i64);
7821            let xv = _mm256_loadu_si256(xq_ptr.add(gi * GROUP_SIZE) as *const __m256i);
7822            let d = dpbusd_hsum(_mm256_abs_epi8(wv), _mm256_sign_epi8(xv, wv));
7823            acc += d as f32 * s;
7824        }
7825        acc
7826    }
7827}
7828
7829/// Per-row int8 base dot, dispatched once per row (matvec decode hot path).
7830/// Callers are gated by `a8w8_enabled()`, so the target-feature kernels are
7831/// reachable.
7832#[inline]
7833fn q1t_dot_row_i8(bytes: &[u8], r: usize, gpr: usize, xq: &[i8]) -> f32 {
7834    #[cfg(target_arch = "aarch64")]
7835    unsafe {
7836        return q1t_dot_row_sdot(bytes, r, gpr, xq);
7837    }
7838    #[cfg(target_arch = "x86_64")]
7839    unsafe {
7840        if vnni_tiles_enabled() {
7841            return q1t_dot_row_vnni(bytes, r, gpr, xq);
7842        }
7843        return q1t_dot_row_avx2(bytes, r, gpr, xq);
7844    }
7845    #[allow(unreachable_code)]
7846    {
7847        const TILE: usize = cortiq_core::quant::Q1T_TILE;
7848        let mut acc = 0f32;
7849        let mut sg = [0i8; GROUP_SIZE + 8]; // +8 slack for the u64-store unpack
7850        for gi in 0..gpr {
7851            let off = (r * gpr + gi) * TILE;
7852            let s = f16_to_f32(u16::from_le_bytes([bytes[off], bytes[off + 1]]));
7853            q1t_unpack_group_i8(bytes.as_ptr().wrapping_add(off + 2), &mut sg);
7854            let mut d = 0i32;
7855            for k in 0..GROUP_SIZE {
7856                d += sg[k] as i32 * xq[gi * GROUP_SIZE + k] as i32;
7857            }
7858            acc += d as f32 * s;
7859        }
7860        acc
7861    }
7862}
7863
7864/// Σ over a row's outliers of `value·x[col]` — the correction that adds the
7865/// overlay's exact weights on top of the base dot. INVARIANT: the encoder
7866/// writes ternary code 0 at every outlier position (`quantize_q1t`), so the
7867/// base contributes nothing there and this is a plain `value·x`, not
7868/// `(value − base)·x` — no scattered per-outlier scale read. Row `r`'s entries
7869/// are the contiguous slice `[row_ptr[r], row_ptr[r+1])`, so no binary search.
7870fn q1t_row_outlier_correction(
7871    bytes: &[u8],
7872    r: usize,
7873    rp_off: usize,
7874    entries_off: usize,
7875    has_ov: bool,
7876    x: &[f32],
7877) -> f32 {
7878    if !has_ov {
7879        return 0.0;
7880    }
7881    let (c0, c1) = (
7882        q1t_rowptr(bytes, rp_off, r),
7883        q1t_rowptr(bytes, rp_off, r + 1),
7884    );
7885    let mut corr = 0f32;
7886    for p in c0..c1 {
7887        let e = entries_off + p * 4;
7888        let col = u16::from_le_bytes([bytes[e], bytes[e + 1]]) as usize;
7889        let val = f16_to_f32(u16::from_le_bytes([bytes[e + 2], bytes[e + 3]]));
7890        corr += val * x[col];
7891    }
7892    corr
7893}
7894
7895/// Dequantize one q1t row into `buf[..cols]` via the sign LUT (no division),
7896/// then apply the row's outliers (its `[row_ptr[r], row_ptr[r+1])` slice).
7897/// Used by the batched (prefill) path where the decode amortizes over the batch.
7898fn q1t_dequant_row(
7899    bytes: &[u8],
7900    r: usize,
7901    gpr: usize,
7902    rp_off: usize,
7903    entries_off: usize,
7904    has_ov: bool,
7905    buf: &mut [f32],
7906) {
7907    const TILE: usize = cortiq_core::quant::Q1T_TILE;
7908    for g in 0..gpr {
7909        let off = (r * gpr + g) * TILE;
7910        let s = f16_to_f32(u16::from_le_bytes([bytes[off], bytes[off + 1]]));
7911        let codes = &bytes[off + 2..off + TILE];
7912        let bc = g * GROUP_SIZE;
7913        // 6 full bytes (30 codes) + a 7th byte holding the last 2.
7914        for bi in 0..6 {
7915            let lut = &SIGN5[codes[bi] as usize];
7916            let d = &mut buf[bc + bi * 5..bc + bi * 5 + 5];
7917            for i in 0..5 {
7918                d[i] = lut[i] * s;
7919            }
7920        }
7921        let lut = &SIGN5[codes[6] as usize];
7922        buf[bc + 30] = lut[0] * s;
7923        buf[bc + 31] = lut[1] * s;
7924    }
7925    if !has_ov {
7926        return;
7927    }
7928    let (c0, c1) = (
7929        q1t_rowptr(bytes, rp_off, r),
7930        q1t_rowptr(bytes, rp_off, r + 1),
7931    );
7932    for p in c0..c1 {
7933        let e = entries_off + p * 4;
7934        let col = u16::from_le_bytes([bytes[e], bytes[e + 1]]) as usize;
7935        buf[col] = f16_to_f32(u16::from_le_bytes([bytes[e + 2], bytes[e + 3]]));
7936    }
7937}
7938
7939/// Add the sparse outlier overlay onto a base dot already in `out` (the GPU
7940/// computes the ternary base; the overlay stays on the CPU — its entries are
7941/// few and its per-row gather doesn't vectorize on the GPU). Row-parallel.
7942fn q1t_add_overlay(
7943    bytes: &[u8],
7944    x: &[f32],
7945    rows: usize,
7946    cols: usize,
7947    out: &mut [f32],
7948    pool: Option<&Pool>,
7949) {
7950    const TILE: usize = cortiq_core::quant::Q1T_TILE;
7951    let gpr = cols / GROUP_SIZE;
7952    let (rp_off, ent_off, has_ov) = q1t_overlay(bytes, rows * gpr * TILE, rows);
7953    if !has_ov {
7954        return;
7955    }
7956    let out_addr = SendMut(out.as_mut_ptr());
7957    let run = move |start: usize, end: usize| {
7958        for r in start..end {
7959            let corr = q1t_row_outlier_correction(bytes, r, rp_off, ent_off, has_ov, x);
7960            // SAFETY: disjoint rows; add onto the base the GPU already wrote.
7961            unsafe { *out_addr.at(r) += corr };
7962        }
7963    };
7964    dispatch_rows(pool, rows, &run);
7965}
7966
7967/// Q1T row range via the A8W8 int8 path — shared activation split,
7968/// per-row: base SDOT dot + outlier correction + overlay.
7969#[allow(clippy::too_many_arguments)]
7970fn q1t_range_a8w8(
7971    bytes: &[u8],
7972    gpr: usize,
7973    rp_off: usize,
7974    ent_off: usize,
7975    has_ov: bool,
7976    act: &SplitAct,
7977    x: &[f32],
7978    out: SendMut,
7979    start: usize,
7980    end: usize,
7981) {
7982    for r in start..end {
7983        let mut acc = q1t_dot_row_i8(bytes, r, gpr, &act.xq) * act.sx;
7984        for &(j, xv) in &act.outliers {
7985            acc += q1t_base_weight(bytes, r, gpr, j) * xv;
7986        }
7987        acc += q1t_row_outlier_correction(bytes, r, rp_off, ent_off, has_ov, x);
7988        // SAFETY: disjoint row ranges per worker.
7989        unsafe { *out.at(r) = acc };
7990    }
7991}
7992
7993/// Q1T row range via the f32 path (no SDOT) — for matvec_many batched
7994/// dispatch when a8w8 is unavailable.
7995#[allow(clippy::too_many_arguments)]
7996fn q1t_range_f32_batch(
7997    bytes: &[u8],
7998    gpr: usize,
7999    rp_off: usize,
8000    ent_off: usize,
8001    has_ov: bool,
8002    x: &[f32],
8003    out: SendMut,
8004    start: usize,
8005    end: usize,
8006) {
8007    const TILE: usize = cortiq_core::quant::Q1T_TILE;
8008    let mut sg = [0f32; GROUP_SIZE];
8009    for r in start..end {
8010        let mut acc = 0f32;
8011        for g in 0..gpr {
8012            let off = (r * gpr + g) * TILE;
8013            let s = f16_to_f32(u16::from_le_bytes([bytes[off], bytes[off + 1]]));
8014            let codes = &bytes[off + 2..off + TILE];
8015            let xg = &x[g * GROUP_SIZE..g * GROUP_SIZE + GROUP_SIZE];
8016            for bi in 0..6 {
8017                sg[bi * 5..bi * 5 + 5].copy_from_slice(&SIGN5[codes[bi] as usize]);
8018            }
8019            let lut = &SIGN5[codes[6] as usize];
8020            sg[30] = lut[0];
8021            sg[31] = lut[1];
8022            let mut gsum = 0f32;
8023            for k in 0..GROUP_SIZE {
8024                gsum += sg[k] * xg[k];
8025            }
8026            acc += s * gsum;
8027        }
8028        acc += q1t_row_outlier_correction(bytes, r, rp_off, ent_off, has_ov, x);
8029        // SAFETY: disjoint row ranges per worker.
8030        unsafe { *out.at(r) = acc };
8031    }
8032}
8033
8034/// Ternary (q1t) matvec — decode+dot straight from mmap, one group at a time:
8035/// no per-ROW buffer, no division (the sign LUT), and a tiny per-group sign
8036/// buffer so the 32-wide dot vectorizes. This is the decode hot path.
8037fn q1t_matvec(
8038    bytes: &[u8],
8039    x: &[f32],
8040    rows: usize,
8041    cols: usize,
8042    out: &mut [f32],
8043    pool: Option<&Pool>,
8044) {
8045    debug_assert_eq!(out.len(), rows);
8046    const TILE: usize = cortiq_core::quant::Q1T_TILE;
8047    let gpr = cols / GROUP_SIZE;
8048    let (rp_off, ent_off, has_ov) = q1t_overlay(bytes, rows * gpr * TILE, rows);
8049    let out_addr = SendMut(out.as_mut_ptr());
8050    // int8 SDOT base dot (ARM dotprod): ~4× the f32 arithmetic. x → i8 once
8051    // (`split_act`), activation outliers added back exactly in f32, weight
8052    // overlay on top. ARM SDOT / x86 AVX2; CMF_SDOT=0 keeps the exact f32 path.
8053    if a8w8_enabled() {
8054        let act = split_act(x);
8055        let act = &act;
8056        let run = move |start: usize, end: usize| {
8057            for r in start..end {
8058                let mut acc = q1t_dot_row_i8(bytes, r, gpr, &act.xq) * act.sx;
8059                for &(j, xv) in &act.outliers {
8060                    acc += q1t_base_weight(bytes, r, gpr, j) * xv;
8061                }
8062                acc += q1t_row_outlier_correction(bytes, r, rp_off, ent_off, has_ov, x);
8063                // SAFETY: disjoint row ranges per worker.
8064                unsafe { *out_addr.at(r) = acc };
8065            }
8066        };
8067        dispatch_rows(pool, rows, &run);
8068        return;
8069    }
8070    let run = move |start: usize, end: usize| {
8071        // Per-group signs, unpacked contiguously so the dot below is a clean
8072        // 32-wide reduction the autovectorizer turns into f32x4 FMAs — the
8073        // 5-values-per-byte base-3 layout won't SIMD in place.
8074        let mut sg = [0f32; GROUP_SIZE];
8075        for r in start..end {
8076            let mut acc = 0f32;
8077            for g in 0..gpr {
8078                let off = (r * gpr + g) * TILE;
8079                let s = f16_to_f32(u16::from_le_bytes([bytes[off], bytes[off + 1]]));
8080                let codes = &bytes[off + 2..off + TILE];
8081                let xg = &x[g * GROUP_SIZE..g * GROUP_SIZE + GROUP_SIZE];
8082                for bi in 0..6 {
8083                    sg[bi * 5..bi * 5 + 5].copy_from_slice(&SIGN5[codes[bi] as usize]);
8084                }
8085                let lut = &SIGN5[codes[6] as usize];
8086                sg[30] = lut[0];
8087                sg[31] = lut[1];
8088                let mut gsum = 0f32;
8089                for k in 0..GROUP_SIZE {
8090                    gsum += sg[k] * xg[k];
8091                }
8092                acc += s * gsum;
8093            }
8094            acc += q1t_row_outlier_correction(bytes, r, rp_off, ent_off, has_ov, x);
8095            unsafe { *out_addr.at(r) = acc };
8096        }
8097    };
8098    dispatch_rows(pool, rows, &run);
8099}
8100
8101/// Fused-pair twin of `q1t_dot_row_sdot`: ONE register unpack of the
8102/// ternary codes serves BOTH activation streams (the unpack chain is
8103/// the dominant per-row cost — MTP verify pairs paid it twice). Per
8104/// stream the group order and f32 accumulation match the single-row
8105/// kernel exactly, so pair == 2×matvec bit-for-bit.
8106#[cfg(target_arch = "aarch64")]
8107#[target_feature(enable = "neon,dotprod")]
8108unsafe fn q1t_dot_row_sdot2(bytes: &[u8], r: usize, gpr: usize, xa: &[i8], xb: &[i8]) -> [f32; 2] {
8109    use core::arch::aarch64::*;
8110    use core::arch::asm;
8111    // SAFETY: same slice-length contracts as `q1t_dot_row_sdot`, ×2.
8112    unsafe {
8113        const TILE: usize = cortiq_core::quant::Q1T_TILE;
8114        let bytes_ptr = bytes.as_ptr();
8115        let row_off = r * gpr * TILE;
8116        let xp = [xa.as_ptr(), xb.as_ptr()];
8117        let mut acc = [0f32; 2];
8118        macro_rules! sdot2 {
8119            ($w0:expr, $w1:expr, $x:expr) => {{
8120                let x0 = vld1q_s8($x);
8121                let x1 = vld1q_s8($x.add(16));
8122                let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
8123                asm!(
8124                    "sdot {a0:v}.4s, {w0:v}.16b, {x0:v}.16b",
8125                    "sdot {a1:v}.4s, {w1:v}.16b, {x1:v}.16b",
8126                    a0 = inout(vreg) a0, a1 = inout(vreg) a1,
8127                    w0 = in(vreg) $w0, x0 = in(vreg) x0, w1 = in(vreg) $w1, x1 = in(vreg) x1,
8128                    options(pure, nomem, nostack),
8129                );
8130                vaddvq_s32(vaddq_s32(a0, a1))
8131            }};
8132        }
8133        let gpr2 = gpr & !1;
8134        let mut gi = 0;
8135        while gi < gpr2 {
8136            let off0 = row_off + gi * TILE;
8137            let off1 = off0 + TILE;
8138            let s0 = f16_to_f32(u16::from_le_bytes([
8139                *bytes_ptr.add(off0),
8140                *bytes_ptr.add(off0 + 1),
8141            ]));
8142            let s1 = f16_to_f32(u16::from_le_bytes([
8143                *bytes_ptr.add(off1),
8144                *bytes_ptr.add(off1 + 1),
8145            ]));
8146            let (u0_0, u1_0, u2_0, u3_0) = q1t_unpack_reg_u64s(bytes_ptr.add(off0 + 2));
8147            let (u0_1, u1_1, u2_1, u3_1) = q1t_unpack_reg_u64s(bytes_ptr.add(off1 + 2));
8148            let w0_0 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u0_0), vcreate_u64(u1_0)));
8149            let w1_0 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u2_0), vcreate_u64(u3_0)));
8150            let w0_1 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u0_1), vcreate_u64(u1_1)));
8151            let w1_1 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u2_1), vcreate_u64(u3_1)));
8152            for k in 0..2 {
8153                let d0 = sdot2!(w0_0, w1_0, xp[k].add(gi * GROUP_SIZE));
8154                let d1 = sdot2!(w0_1, w1_1, xp[k].add((gi + 1) * GROUP_SIZE));
8155                acc[k] += d0 as f32 * s0 + d1 as f32 * s1;
8156            }
8157            gi += 2;
8158        }
8159        if gi < gpr {
8160            let off = row_off + gi * TILE;
8161            let s = f16_to_f32(u16::from_le_bytes([
8162                *bytes_ptr.add(off),
8163                *bytes_ptr.add(off + 1),
8164            ]));
8165            let (u0, u1, u2, u3) = q1t_unpack_reg_u64s(bytes_ptr.add(off + 2));
8166            let w0 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u0), vcreate_u64(u1)));
8167            let w1 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u2), vcreate_u64(u3)));
8168            for k in 0..2 {
8169                let d = sdot2!(w0, w1, xp[k].add(gi * GROUP_SIZE));
8170                acc[k] += d as f32 * s;
8171            }
8172        }
8173        acc
8174    }
8175}
8176
8177/// Fused Q1T pair matvec: ONE pass over the rows serves both
8178/// activation streams — on ARM the ternary register unpack happens
8179/// once per tile pair (`q1t_dot_row_sdot2`); elsewhere the second dot
8180/// rides the row's L1-warm tile bytes. Per stream the math matches
8181/// `q1t_matvec` exactly.
8182fn q1t_matvec2(
8183    bytes: &[u8],
8184    x1: &[f32],
8185    x2: &[f32],
8186    rows: usize,
8187    cols: usize,
8188    o1: &mut [f32],
8189    o2: &mut [f32],
8190    pool: Option<&Pool>,
8191) {
8192    debug_assert_eq!(o1.len(), rows);
8193    debug_assert_eq!(o2.len(), rows);
8194    const TILE: usize = cortiq_core::quant::Q1T_TILE;
8195    let gpr = cols / GROUP_SIZE;
8196    let (rp_off, ent_off, has_ov) = q1t_overlay(bytes, rows * gpr * TILE, rows);
8197    let out1 = SendMut(o1.as_mut_ptr());
8198    let out2 = SendMut(o2.as_mut_ptr());
8199    if a8w8_enabled() {
8200        let a1 = split_act(x1);
8201        let a2 = split_act(x2);
8202        let (a1, a2) = (&a1, &a2);
8203        let run = move |start: usize, end: usize| {
8204            for r in start..end {
8205                #[cfg(target_arch = "aarch64")]
8206                // a8w8 on aarch64 ⇔ sdot_enabled(), so the kernel's
8207                // target features are present.
8208                let ds = unsafe { q1t_dot_row_sdot2(bytes, r, gpr, &a1.xq, &a2.xq) };
8209                #[cfg(not(target_arch = "aarch64"))]
8210                let ds = [
8211                    q1t_dot_row_i8(bytes, r, gpr, &a1.xq),
8212                    q1t_dot_row_i8(bytes, r, gpr, &a2.xq),
8213                ];
8214                let mut acc1 = ds[0] * a1.sx;
8215                for &(j, xv) in &a1.outliers {
8216                    acc1 += q1t_base_weight(bytes, r, gpr, j) * xv;
8217                }
8218                acc1 += q1t_row_outlier_correction(bytes, r, rp_off, ent_off, has_ov, x1);
8219                let mut acc2 = ds[1] * a2.sx;
8220                for &(j, xv) in &a2.outliers {
8221                    acc2 += q1t_base_weight(bytes, r, gpr, j) * xv;
8222                }
8223                acc2 += q1t_row_outlier_correction(bytes, r, rp_off, ent_off, has_ov, x2);
8224                // SAFETY: disjoint row ranges per worker.
8225                unsafe {
8226                    *out1.at(r) = acc1;
8227                    *out2.at(r) = acc2;
8228                }
8229            }
8230        };
8231        dispatch_rows(pool, rows, &run);
8232        return;
8233    }
8234    let run = move |start: usize, end: usize| {
8235        // Exact path (CMF_SDOT=0): unpack the sign LUT once per group,
8236        // dot both streams — same op order per stream as `q1t_matvec`.
8237        let mut sg = [0f32; GROUP_SIZE];
8238        for r in start..end {
8239            let mut acc1 = 0f32;
8240            let mut acc2 = 0f32;
8241            for g in 0..gpr {
8242                let off = (r * gpr + g) * TILE;
8243                let s = f16_to_f32(u16::from_le_bytes([bytes[off], bytes[off + 1]]));
8244                let codes = &bytes[off + 2..off + TILE];
8245                for bi in 0..6 {
8246                    sg[bi * 5..bi * 5 + 5].copy_from_slice(&SIGN5[codes[bi] as usize]);
8247                }
8248                let lut = &SIGN5[codes[6] as usize];
8249                sg[30] = lut[0];
8250                sg[31] = lut[1];
8251                let xg1 = &x1[g * GROUP_SIZE..g * GROUP_SIZE + GROUP_SIZE];
8252                let xg2 = &x2[g * GROUP_SIZE..g * GROUP_SIZE + GROUP_SIZE];
8253                let mut gsum1 = 0f32;
8254                for k in 0..GROUP_SIZE {
8255                    gsum1 += sg[k] * xg1[k];
8256                }
8257                acc1 += s * gsum1;
8258                let mut gsum2 = 0f32;
8259                for k in 0..GROUP_SIZE {
8260                    gsum2 += sg[k] * xg2[k];
8261                }
8262                acc2 += s * gsum2;
8263            }
8264            acc1 += q1t_row_outlier_correction(bytes, r, rp_off, ent_off, has_ov, x1);
8265            acc2 += q1t_row_outlier_correction(bytes, r, rp_off, ent_off, has_ov, x2);
8266            // SAFETY: disjoint row ranges per worker.
8267            unsafe {
8268                *out1.at(r) = acc1;
8269                *out2.at(r) = acc2;
8270            }
8271        }
8272    };
8273    dispatch_rows(pool, rows, &run);
8274}
8275
8276/// Ternary (q1t) matmat (prefill) — dequant each row once, dot the whole
8277/// batch against it (amortizes the per-row decode).
8278fn q1t_matmat(
8279    bytes: &[u8],
8280    xs: &[f32],
8281    b: usize,
8282    rows: usize,
8283    cols: usize,
8284    out: &mut [f32],
8285    pool: Option<&Pool>,
8286) {
8287    debug_assert_eq!(out.len(), b * rows);
8288    const TILE: usize = cortiq_core::quant::Q1T_TILE;
8289    let gpr = cols / GROUP_SIZE;
8290    let (rp_off, ent_off, has_ov) = q1t_overlay(bytes, rows * gpr * TILE, rows);
8291    let out_addr = SendMut(out.as_mut_ptr());
8292    // int8 prefill (ARM SDOT / x86 AVX2): quantize the B inputs once, unpack
8293    // each weight row's signs to i8 ONCE, then int8-dot against every input —
8294    // the row sign-decode amortizes over the whole batch. CMF_SDOT=0 → f32.
8295    if a8w8_enabled() {
8296        let acts: Vec<SplitAct> = (0..b)
8297            .map(|bi| split_act(&xs[bi * cols..(bi + 1) * cols]))
8298            .collect();
8299        let acts = &acts;
8300        let run = move |start: usize, end: usize| {
8301            let mut sg = vec![0i8; cols + 8]; // row signs, i8 (+8 unpack slack)
8302            let mut sc = vec![0f32; gpr]; // per-group scales
8303            let mut accs = vec![0f32; b]; // per-batch accumulators, reused per row
8304            for r in start..end {
8305                for g in 0..gpr {
8306                    let off = (r * gpr + g) * TILE;
8307                    sc[g] = f16_to_f32(u16::from_le_bytes([bytes[off], bytes[off + 1]]));
8308                    q1t_unpack_group_i8(
8309                        bytes.as_ptr().wrapping_add(off + 2),
8310                        &mut sg[g * GROUP_SIZE..],
8311                    );
8312                }
8313                for bi in 0..b {
8314                    let act = &acts[bi];
8315                    let mut isum = 0f32;
8316                    for g in 0..gpr {
8317                        let d = q1t_i8dot32(
8318                            sg.as_ptr().wrapping_add(g * GROUP_SIZE),
8319                            act.xq.as_ptr().wrapping_add(g * GROUP_SIZE),
8320                        );
8321                        isum += d as f32 * sc[g];
8322                    }
8323                    let mut acc = isum * act.sx;
8324                    for &(j, xv) in &act.outliers {
8325                        acc += q1t_base_weight(bytes, r, gpr, j) * xv;
8326                    }
8327                    accs[bi] = acc;
8328                }
8329                // Overlay ONCE per row for the whole batch: read each (col, val)
8330                // from mmap a single time (was b× — the re-read dominated prefill)
8331                // and fan it out over the batch via the cached inputs.
8332                if has_ov {
8333                    let (c0, c1) = (
8334                        q1t_rowptr(bytes, rp_off, r),
8335                        q1t_rowptr(bytes, rp_off, r + 1),
8336                    );
8337                    for p in c0..c1 {
8338                        let e = ent_off + p * 4;
8339                        let col = u16::from_le_bytes([bytes[e], bytes[e + 1]]) as usize;
8340                        let val = f16_to_f32(u16::from_le_bytes([bytes[e + 2], bytes[e + 3]]));
8341                        for bi in 0..b {
8342                            accs[bi] += val * xs[bi * cols + col];
8343                        }
8344                    }
8345                }
8346                for bi in 0..b {
8347                    unsafe { *out_addr.at(bi * rows + r) = accs[bi] };
8348                }
8349            }
8350        };
8351        dispatch_rows(pool, rows, &run);
8352        return;
8353    }
8354    let run = move |start: usize, end: usize| {
8355        let mut buf = vec![0f32; cols];
8356        for r in start..end {
8357            q1t_dequant_row(bytes, r, gpr, rp_off, ent_off, has_ov, &mut buf);
8358            for bi in 0..b {
8359                let xr = &xs[bi * cols..(bi + 1) * cols];
8360                let mut acc = 0f32;
8361                for j in 0..cols {
8362                    acc += buf[j] * xr[j];
8363                }
8364                unsafe { *out_addr.at(bi * rows + r) = acc };
8365            }
8366        }
8367    };
8368    dispatch_rows(pool, rows, &run);
8369}
8370
8371fn q1_matvec(
8372    bytes: &[u8],
8373    x: &[f32],
8374    rows: usize,
8375    cols: usize,
8376    out: &mut [f32],
8377    pool: Option<&Pool>,
8378) {
8379    debug_assert_eq!(out.len(), rows);
8380    let gpr = cols / GROUP_SIZE;
8381    let out_addr = SendMut(out.as_mut_ptr());
8382    if a8w8_enabled() {
8383        let act = split_act(x);
8384        let gsum = q1_group_sums(&act.xq, gpr);
8385        let (act, gsum) = (&act, &gsum);
8386        let run = move |start: usize, end: usize| {
8387            q1_range_a8w8(bytes, gpr, act, gsum, out_addr, start, end)
8388        };
8389        dispatch_rows(pool, rows, &run);
8390        return;
8391    }
8392    let run = move |start: usize, end: usize| q1_range_f32(bytes, gpr, x, out_addr, start, end);
8393    dispatch_rows(pool, rows, &run);
8394}
8395
8396/// Fused two-input q1 matvec (weights read once per pair).
8397#[allow(clippy::too_many_arguments)]
8398fn q1_matvec2(
8399    bytes: &[u8],
8400    x1: &[f32],
8401    x2: &[f32],
8402    rows: usize,
8403    cols: usize,
8404    o1: &mut [f32],
8405    o2: &mut [f32],
8406    pool: Option<&Pool>,
8407) {
8408    let gpr = cols / GROUP_SIZE;
8409    let p1 = SendMut(o1.as_mut_ptr());
8410    let p2 = SendMut(o2.as_mut_ptr());
8411    if a8w8_enabled() {
8412        let a1 = split_act(x1);
8413        let a2 = split_act(x2);
8414        let g1 = q1_group_sums(&a1.xq, gpr);
8415        let g2 = q1_group_sums(&a2.xq, gpr);
8416        let (a1, a2, g1, g2) = (&a1, &a2, &g1, &g2);
8417        let run = move |start: usize, end: usize| {
8418            for r in start..end {
8419                let mut v1 = dot_q1_row_i8(bytes, r, gpr, &a1.xq, g1) * a1.sx;
8420                let mut v2 = dot_q1_row_i8(bytes, r, gpr, &a2.xq, g2) * a2.sx;
8421                for &(j, xv) in &a1.outliers {
8422                    let (w, s) = q1_outlier(bytes, r, gpr, j);
8423                    v1 += w * s * xv;
8424                }
8425                for &(j, xv) in &a2.outliers {
8426                    let (w, s) = q1_outlier(bytes, r, gpr, j);
8427                    v2 += w * s * xv;
8428                }
8429                // SAFETY: disjoint row ranges per worker.
8430                unsafe {
8431                    *p1.at(r) = v1;
8432                    *p2.at(r) = v2;
8433                }
8434            }
8435        };
8436        dispatch_rows(pool, rows, &run);
8437        return;
8438    }
8439    let run = move |start: usize, end: usize| {
8440        for r in start..end {
8441            // SAFETY: disjoint row ranges per worker.
8442            unsafe {
8443                *p1.at(r) = q1_row_exact(bytes, r, gpr, x1);
8444                *p2.at(r) = q1_row_exact(bytes, r, gpr, x2);
8445            }
8446        }
8447    };
8448    dispatch_rows(pool, rows, &run);
8449}
8450
8451/// Batched q1 matmat: each row's tiles stream once per microbatch.
8452#[allow(clippy::too_many_arguments)]
8453fn q1_matmat(
8454    bytes: &[u8],
8455    xs_all: &[f32],
8456    b: usize,
8457    rows: usize,
8458    cols: usize,
8459    out: &mut [f32],
8460    pool: Option<&Pool>,
8461) {
8462    debug_assert_eq!(out.len(), b * rows);
8463    let gpr = cols / GROUP_SIZE;
8464    let out_addr = SendMut(out.as_mut_ptr());
8465    if a8w8_enabled() {
8466        let acts: Vec<(SplitAct, Vec<i32>)> = (0..b)
8467            .map(|bi| {
8468                let act = split_act(&xs_all[bi * cols..(bi + 1) * cols]);
8469                let gsum = q1_group_sums(&act.xq, gpr);
8470                (act, gsum)
8471            })
8472            .collect();
8473        let acts = &acts;
8474        #[cfg(target_arch = "x86_64")]
8475        let blocked_ok = avx2_enabled() && blocked_enabled();
8476        #[cfg(target_arch = "aarch64")]
8477        let blocked_ok = sdot_enabled() && blocked_enabled();
8478        let run = move |start: usize, end: usize| {
8479            for r in start..end {
8480                let mut bi = 0usize;
8481                // Blocked 1×4: the unpacked bit mask serves four
8482                // activation streams per group.
8483                #[cfg(target_arch = "aarch64")]
8484                if blocked_ok {
8485                    while bi + 4 <= acts.len() {
8486                        let xs = [
8487                            acts[bi].0.xq.as_slice(),
8488                            acts[bi + 1].0.xq.as_slice(),
8489                            acts[bi + 2].0.xq.as_slice(),
8490                            acts[bi + 3].0.xq.as_slice(),
8491                        ];
8492                        let gs = [
8493                            acts[bi].1.as_slice(),
8494                            acts[bi + 1].1.as_slice(),
8495                            acts[bi + 2].1.as_slice(),
8496                            acts[bi + 3].1.as_slice(),
8497                        ];
8498                        let d = unsafe { dot_q1_row_1x4_sdot(bytes, r, gpr, xs, gs) };
8499                        for k in 0..4 {
8500                            let (act, _) = &acts[bi + k];
8501                            let mut acc = d[k] * act.sx;
8502                            for &(j, xv) in &act.outliers {
8503                                let (w, sc) = q1_outlier(bytes, r, gpr, j);
8504                                acc += w * sc * xv;
8505                            }
8506                            // SAFETY: disjoint (bi, r) cells per worker.
8507                            unsafe { *out_addr.at((bi + k) * rows + r) = acc };
8508                        }
8509                        bi += 4;
8510                    }
8511                }
8512                #[cfg(target_arch = "x86_64")]
8513                if blocked_ok {
8514                    while bi + 4 <= acts.len() {
8515                        let xs = [
8516                            acts[bi].0.xq.as_slice(),
8517                            acts[bi + 1].0.xq.as_slice(),
8518                            acts[bi + 2].0.xq.as_slice(),
8519                            acts[bi + 3].0.xq.as_slice(),
8520                        ];
8521                        let gs = [
8522                            acts[bi].1.as_slice(),
8523                            acts[bi + 1].1.as_slice(),
8524                            acts[bi + 2].1.as_slice(),
8525                            acts[bi + 3].1.as_slice(),
8526                        ];
8527                        let d = unsafe {
8528                            if vnni_tiles_enabled() {
8529                                dot_q1_row_1x4_vnni(bytes, r, gpr, xs, gs)
8530                            } else {
8531                                dot_q1_row_1x4_avx2(bytes, r, gpr, xs, gs)
8532                            }
8533                        };
8534                        for k in 0..4 {
8535                            let (act, _) = &acts[bi + k];
8536                            let mut acc = d[k] * act.sx;
8537                            for &(j, xv) in &act.outliers {
8538                                let (w, sc) = q1_outlier(bytes, r, gpr, j);
8539                                acc += w * sc * xv;
8540                            }
8541                            // SAFETY: disjoint (bi, r) cells per worker.
8542                            unsafe { *out_addr.at((bi + k) * rows + r) = acc };
8543                        }
8544                        bi += 4;
8545                    }
8546                }
8547                while bi < acts.len() {
8548                    let (act, gsum) = &acts[bi];
8549                    let mut acc = dot_q1_row_i8(bytes, r, gpr, &act.xq, gsum) * act.sx;
8550                    for &(j, xv) in &act.outliers {
8551                        let (w, s) = q1_outlier(bytes, r, gpr, j);
8552                        acc += w * s * xv;
8553                    }
8554                    // SAFETY: disjoint (bi, r) cells per worker range.
8555                    unsafe { *out_addr.at(bi * rows + r) = acc };
8556                    bi += 1;
8557                }
8558            }
8559        };
8560        dispatch_rows(pool, rows, &run);
8561        return;
8562    }
8563    let run = move |start: usize, end: usize| {
8564        for r in start..end {
8565            for bi in 0..b {
8566                let x = &xs_all[bi * cols..(bi + 1) * cols];
8567                // SAFETY: disjoint (bi, r) cells per worker range.
8568                unsafe { *out_addr.at(bi * rows + r) = q1_row_exact(bytes, r, gpr, x) };
8569            }
8570        }
8571    };
8572    dispatch_rows(pool, rows, &run);
8573}
8574
8575/// Fused q4_block matvec straight from the mapped bytes. SDOT path when
8576/// dotprod is available (port of vmfcore `dot_q4_block_sdot`, measured
8577/// +23% on q4 decode): nibbles → centered i8, int8×int8 `sdot` per
8578/// 32-group, exact outlier correction — the same A8W8 contract as q8.
8579/// `CMF_SDOT=0` keeps the exact scalar path.
8580fn q4matvec(
8581    bytes: &[u8],
8582    x: &[f32],
8583    rows: usize,
8584    cols: usize,
8585    out: &mut [f32],
8586    pool: Option<&Pool>,
8587) {
8588    debug_assert_eq!(out.len(), rows);
8589    let (packed, scales) = q4_split(bytes, rows, cols);
8590    let gpr = cols / GROUP_SIZE;
8591    let out_addr = SendMut(out.as_mut_ptr());
8592
8593    if a8w8_enabled() {
8594        let act = split_act(x);
8595        let run = move |start: usize, end: usize| {
8596            q4_range_a8w8(packed, scales, gpr, cols, &act, out_addr, start, end)
8597        };
8598        dispatch_rows(pool, rows, &run);
8599        return;
8600    }
8601
8602    let run =
8603        move |start: usize, end: usize| q4_range_f32(packed, scales, gpr, x, out_addr, start, end);
8604    dispatch_rows(pool, rows, &run);
8605}
8606
8607/// One q4 row via the A8W8 int8 path — SDOT on ARM, AVX2 maddubs on
8608/// x86 (scalar fallback is unreachable: callers gate on a8w8_enabled).
8609#[inline]
8610#[allow(unreachable_code)]
8611/// One UNPACKED q4 row (centered i8 in `buf`) against four activation
8612/// streams: the 32-byte weight chunk and its abs() load once per group,
8613/// the per-group f16 scale decodes once — four maddubs+reduce chains
8614/// instead of four full (load, abs, dot) rounds.
8615#[cfg(target_arch = "x86_64")]
8616#[target_feature(enable = "avx2")]
8617unsafe fn dot_q4b_row_1x4_avx2(
8618    buf: &[u8],
8619    scales: &[u8],
8620    g0: usize,
8621    gpr: usize,
8622    xs: [&[i8]; 4],
8623) -> [f32; 4] {
8624    // SAFETY: callers uphold buffer contracts (buf.len() == gpr·32).
8625    unsafe {
8626        use core::arch::x86_64::*;
8627        let ones = _mm256_set1_epi16(1);
8628        let mut acc = [0f32; 4];
8629        for gi in 0..gpr {
8630            let s = f16_to_f32(u16::from_le_bytes([
8631                scales[(g0 + gi) * 2],
8632                scales[(g0 + gi) * 2 + 1],
8633            ]));
8634            let w = _mm256_loadu_si256(buf.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
8635            let aw = _mm256_abs_epi8(w);
8636            for (k, xq) in xs.iter().enumerate() {
8637                let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
8638                let p16 = _mm256_maddubs_epi16(aw, _mm256_sign_epi8(x, w));
8639                let d = _mm256_madd_epi16(p16, ones);
8640                let hi128 = _mm256_extracti128_si256::<1>(d);
8641                let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
8642                let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
8643                let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
8644                acc[k] += _mm_cvtsi128_si32(s32) as f32 * s;
8645            }
8646        }
8647        acc
8648    }
8649}
8650
8651/// VNNI twin of `dot_q4b_row_1x4_avx2` (see `dpbusd_hsum`).
8652#[cfg(target_arch = "x86_64")]
8653#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
8654unsafe fn dot_q4b_row_1x4_vnni(
8655    buf: &[u8],
8656    scales: &[u8],
8657    g0: usize,
8658    gpr: usize,
8659    xs: [&[i8]; 4],
8660) -> [f32; 4] {
8661    // SAFETY: callers uphold buffer contracts (buf.len() == gpr·32).
8662    unsafe {
8663        use core::arch::x86_64::*;
8664        let mut acc = [0f32; 4];
8665        for gi in 0..gpr {
8666            let s = f16_to_f32(u16::from_le_bytes([
8667                scales[(g0 + gi) * 2],
8668                scales[(g0 + gi) * 2 + 1],
8669            ]));
8670            let w = _mm256_loadu_si256(buf.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
8671            let aw = _mm256_abs_epi8(w);
8672            for (k, xq) in xs.iter().enumerate() {
8673                let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
8674                let d = dpbusd_hsum(aw, _mm256_sign_epi8(x, w));
8675                acc[k] += d as f32 * s;
8676            }
8677        }
8678        acc
8679    }
8680}
8681
8682/// The vbit flavor of the blocked 1×4: the per-activation A8W8 scale
8683/// folds in PER GROUP as `(d·sx)·s` — bit-matching the single-matvec
8684/// accumulation order (the q4_block flavor applies sx once at the end,
8685/// matching ITS single path; the two conventions are historical and
8686/// each blocked leg must mirror its own).
8687#[cfg(target_arch = "x86_64")]
8688#[target_feature(enable = "avx2")]
8689unsafe fn dot_q4b_row_1x4_sx_avx2(
8690    buf: &[u8],
8691    scales: &[u8],
8692    g0: usize,
8693    gpr: usize,
8694    xs: [&[i8]; 4],
8695    sxs: [f32; 4],
8696) -> [f32; 4] {
8697    // SAFETY: callers uphold buffer contracts (buf.len() == gpr·32).
8698    unsafe {
8699        use core::arch::x86_64::*;
8700        let ones = _mm256_set1_epi16(1);
8701        let mut acc = [0f32; 4];
8702        for gi in 0..gpr {
8703            let s = f16_to_f32(u16::from_le_bytes([
8704                scales[(g0 + gi) * 2],
8705                scales[(g0 + gi) * 2 + 1],
8706            ]));
8707            let w = _mm256_loadu_si256(buf.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
8708            let aw = _mm256_abs_epi8(w);
8709            for (k, xq) in xs.iter().enumerate() {
8710                let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
8711                let p16 = _mm256_maddubs_epi16(aw, _mm256_sign_epi8(x, w));
8712                let d = _mm256_madd_epi16(p16, ones);
8713                let hi128 = _mm256_extracti128_si256::<1>(d);
8714                let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
8715                let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
8716                let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
8717                acc[k] += (_mm_cvtsi128_si32(s32) as f32 * sxs[k]) * s;
8718            }
8719        }
8720        acc
8721    }
8722}
8723
8724/// VNNI twin of `dot_q4b_row_1x4_sx_avx2` (see `dpbusd_hsum`; the
8725/// per-group `(d·sx)·s` fold mirrors the vbit single path).
8726#[cfg(target_arch = "x86_64")]
8727#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
8728unsafe fn dot_q4b_row_1x4_sx_vnni(
8729    buf: &[u8],
8730    scales: &[u8],
8731    g0: usize,
8732    gpr: usize,
8733    xs: [&[i8]; 4],
8734    sxs: [f32; 4],
8735) -> [f32; 4] {
8736    // SAFETY: callers uphold buffer contracts (buf.len() == gpr·32).
8737    unsafe {
8738        use core::arch::x86_64::*;
8739        let mut acc = [0f32; 4];
8740        for gi in 0..gpr {
8741            let s = f16_to_f32(u16::from_le_bytes([
8742                scales[(g0 + gi) * 2],
8743                scales[(g0 + gi) * 2 + 1],
8744            ]));
8745            let w = _mm256_loadu_si256(buf.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
8746            let aw = _mm256_abs_epi8(w);
8747            for (k, xq) in xs.iter().enumerate() {
8748                let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
8749                let d = dpbusd_hsum(aw, _mm256_sign_epi8(x, w));
8750                acc[k] += (d as f32 * sxs[k]) * s;
8751            }
8752        }
8753        acc
8754    }
8755}
8756
8757#[allow(unreachable_code)]
8758fn dot_q4_row_i8(packed: &[u8], scales: &[u8], g0: usize, gpr: usize, xq: &[i8]) -> f32 {
8759    #[cfg(target_arch = "aarch64")]
8760    unsafe {
8761        return dot_q4_row_sdot(packed, scales, g0, gpr, xq);
8762    }
8763    #[cfg(target_arch = "x86_64")]
8764    unsafe {
8765        return dot_q4_row_avx2(packed, scales, g0, gpr, xq);
8766    }
8767    let mut acc = 0f32;
8768    for gi in 0..gpr {
8769        let g = g0 + gi;
8770        let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
8771        let mut d = 0i32;
8772        for (k, &b) in packed[g * 16..(g + 1) * 16].iter().enumerate() {
8773            d += ((b & 0x0F) as i32 - 8) * xq[gi * GROUP_SIZE + k * 2] as i32
8774                + (((b >> 4) & 0x0F) as i32 - 8) * xq[gi * GROUP_SIZE + k * 2 + 1] as i32;
8775        }
8776        acc += d as f32 * s;
8777    }
8778    acc
8779}
8780
8781/// Two-activation q4 row via the A8W8 int8 path (see `dot_q4_row_i8`).
8782#[inline]
8783#[allow(unreachable_code)]
8784fn dot_q4_row_i8_2(
8785    packed: &[u8],
8786    scales: &[u8],
8787    g0: usize,
8788    gpr: usize,
8789    xq1: &[i8],
8790    xq2: &[i8],
8791) -> (f32, f32) {
8792    #[cfg(target_arch = "aarch64")]
8793    unsafe {
8794        return dot_q4_row_sdot2(packed, scales, g0, gpr, xq1, xq2);
8795    }
8796    #[cfg(target_arch = "x86_64")]
8797    unsafe {
8798        return dot_q4_row_avx2_2(packed, scales, g0, gpr, xq1, xq2);
8799    }
8800    (
8801        dot_q4_row_i8(packed, scales, g0, gpr, xq1),
8802        dot_q4_row_i8(packed, scales, g0, gpr, xq2),
8803    )
8804}
8805
8806/// One q4 row range via SDOT (kernel body of `q4matvec`, extracted so
8807/// multi-matrix jobs can drive it for several tensors in one dispatch).
8808#[allow(clippy::too_many_arguments)]
8809fn q4_range_a8w8(
8810    packed: &[u8],
8811    scales: &[u8],
8812    gpr: usize,
8813    cols: usize,
8814    act: &SplitAct,
8815    out: SendMut,
8816    start: usize,
8817    end: usize,
8818) {
8819    for r in start..end {
8820        let mut acc = dot_q4_row_i8(packed, scales, r * gpr, gpr, &act.xq) * act.sx;
8821        // xq is zeroed at outlier slots — add the exact terms.
8822        for &(j, xv) in &act.outliers {
8823            let flat = r * cols + j;
8824            let byte = packed[flat / 2];
8825            let nib = if flat & 1 == 0 {
8826                byte & 0x0F
8827            } else {
8828                byte >> 4
8829            };
8830            let s = f16_to_f32(u16::from_le_bytes([
8831                scales[(flat / GROUP_SIZE) * 2],
8832                scales[(flat / GROUP_SIZE) * 2 + 1],
8833            ]));
8834            acc += ((nib as i32 - 8) as f32) * s * xv;
8835        }
8836        // SAFETY: disjoint row ranges per worker.
8837        unsafe { *out.at(r) = acc };
8838    }
8839}
8840
8841/// Two-input q4 row range via the A8W8 int8 path — kernel body of
8842/// `q4matvec2`, extracted for pair multi-matrix jobs.
8843#[allow(clippy::too_many_arguments)]
8844fn q4_range2_a8w8(
8845    packed: &[u8],
8846    scales: &[u8],
8847    gpr: usize,
8848    cols: usize,
8849    a1: &SplitAct,
8850    a2: &SplitAct,
8851    p1: SendMut,
8852    p2: SendMut,
8853    start: usize,
8854    end: usize,
8855) {
8856    for r in start..end {
8857        let (s1, s2) = dot_q4_row_i8_2(packed, scales, r * gpr, gpr, &a1.xq, &a2.xq);
8858        let mut acc1 = s1 * a1.sx;
8859        let mut acc2 = s2 * a2.sx;
8860        // xq is zeroed at outlier slots — add the exact terms.
8861        let fix = |outliers: &[(usize, f32)], acc: &mut f32| {
8862            for &(j, xv) in outliers {
8863                let flat = r * cols + j;
8864                let byte = packed[flat / 2];
8865                let nib = if flat & 1 == 0 {
8866                    byte & 0x0F
8867                } else {
8868                    byte >> 4
8869                };
8870                let s = f16_to_f32(u16::from_le_bytes([
8871                    scales[(flat / GROUP_SIZE) * 2],
8872                    scales[(flat / GROUP_SIZE) * 2 + 1],
8873                ]));
8874                *acc += ((nib as i32 - 8) as f32) * s * xv;
8875            }
8876        };
8877        fix(&a1.outliers, &mut acc1);
8878        fix(&a2.outliers, &mut acc2);
8879        // SAFETY: disjoint row ranges per worker.
8880        unsafe {
8881            *p1.at(r) = acc1;
8882            *p2.at(r) = acc2;
8883        }
8884    }
8885}
8886
8887/// Exact scalar q4 row range (same extraction, non-SDOT path).
8888fn q4_range_f32(
8889    packed: &[u8],
8890    scales: &[u8],
8891    gpr: usize,
8892    x: &[f32],
8893    out: SendMut,
8894    start: usize,
8895    end: usize,
8896) {
8897    for r in start..end {
8898        let mut acc = 0f32;
8899        for gi in 0..gpr {
8900            let g = r * gpr + gi;
8901            let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
8902            let pk = &packed[g * 16..(g + 1) * 16];
8903            let xg = &x[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
8904            let mut ga = 0f32;
8905            for (k, &b) in pk.iter().enumerate() {
8906                ga += ((b & 0x0F) as f32 - 8.0) * xg[k * 2]
8907                    + (((b >> 4) & 0x0F) as f32 - 8.0) * xg[k * 2 + 1];
8908            }
8909            acc += ga * s;
8910        }
8911        // SAFETY: disjoint row ranges per worker.
8912        unsafe { *out.at(r) = acc };
8913    }
8914}
8915
8916/// Fused two-input q4 matvec: nibbles are unpacked ONCE per group and
8917/// dotted against both activations (was: two full matvecs — double
8918/// weight traffic). Per-lane math matches `q4matvec` exactly.
8919#[allow(clippy::too_many_arguments)]
8920fn q4matvec2(
8921    bytes: &[u8],
8922    x1: &[f32],
8923    x2: &[f32],
8924    rows: usize,
8925    cols: usize,
8926    o1: &mut [f32],
8927    o2: &mut [f32],
8928    pool: Option<&Pool>,
8929) {
8930    debug_assert_eq!(o1.len(), rows);
8931    debug_assert_eq!(o2.len(), rows);
8932    let (packed, scales) = q4_split(bytes, rows, cols);
8933    let gpr = cols / GROUP_SIZE;
8934
8935    if a8w8_enabled() {
8936        let a1 = split_act(x1);
8937        let a2 = split_act(x2);
8938        let p1 = SendMut(o1.as_mut_ptr());
8939        let p2 = SendMut(o2.as_mut_ptr());
8940        let run = move |start: usize, end: usize| {
8941            q4_range2_a8w8(packed, scales, gpr, cols, &a1, &a2, p1, p2, start, end)
8942        };
8943        dispatch_rows(pool, rows, &run);
8944        return;
8945    }
8946
8947    let p1 = SendMut(o1.as_mut_ptr());
8948    let p2 = SendMut(o2.as_mut_ptr());
8949    let run = move |start: usize, end: usize| {
8950        q4_range2_f32(packed, scales, gpr, x1, x2, p1, p2, start, end)
8951    };
8952    dispatch_rows(pool, rows, &run);
8953}
8954
8955/// Two-input exact scalar q4 row range (same extraction).
8956#[allow(clippy::too_many_arguments)]
8957fn q4_range2_f32(
8958    packed: &[u8],
8959    scales: &[u8],
8960    gpr: usize,
8961    x1: &[f32],
8962    x2: &[f32],
8963    p1: SendMut,
8964    p2: SendMut,
8965    start: usize,
8966    end: usize,
8967) {
8968    for r in start..end {
8969        let (mut acc1, mut acc2) = (0f32, 0f32);
8970        for gi in 0..gpr {
8971            let g = r * gpr + gi;
8972            let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
8973            let pk = &packed[g * 16..(g + 1) * 16];
8974            let x1g = &x1[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
8975            let x2g = &x2[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
8976            let (mut g1, mut g2) = (0f32, 0f32);
8977            for (k, &b) in pk.iter().enumerate() {
8978                let wl = (b & 0x0F) as f32 - 8.0;
8979                let wh = ((b >> 4) & 0x0F) as f32 - 8.0;
8980                g1 += wl * x1g[k * 2] + wh * x1g[k * 2 + 1];
8981                g2 += wl * x2g[k * 2] + wh * x2g[k * 2 + 1];
8982            }
8983            acc1 += g1 * s;
8984            acc2 += g2 * s;
8985        }
8986        // SAFETY: disjoint row ranges per worker.
8987        unsafe {
8988            *p1.at(r) = acc1;
8989            *p2.at(r) = acc2;
8990        }
8991    }
8992}
8993
8994thread_local! {
8995    /// Per-worker decoded-row scratch for the batched q4/vbit kernels
8996    /// (centered i8 for SDOT, f32 for the exact/scalar paths).
8997    static ROW_I8: std::cell::RefCell<Vec<u8>> = const { std::cell::RefCell::new(Vec::new()) };
8998    static ROW_F32: std::cell::RefCell<Vec<f32>> = const { std::cell::RefCell::new(Vec::new()) };
8999}
9000
9001/// Batched q4 matmat: each weight row is unpacked from the mmap ONCE
9002/// and dotted against ALL b activations (prefill used to fall back to b
9003/// full matvecs — b× weight traffic and b× nibble decode). Per-position
9004/// math matches `q4matvec` exactly: same group order, same accumulation.
9005/// `out` is row-major [b, rows] like `qmatmat`.
9006#[allow(clippy::too_many_arguments)]
9007fn q4matmat(
9008    bytes: &[u8],
9009    xs_all: &[f32],
9010    b: usize,
9011    rows: usize,
9012    cols: usize,
9013    out: &mut [f32],
9014    pool: Option<&Pool>,
9015) {
9016    debug_assert_eq!(xs_all.len(), b * cols);
9017    debug_assert_eq!(out.len(), b * rows);
9018    let (packed, scales) = q4_split(bytes, rows, cols);
9019    let gpr = cols / GROUP_SIZE;
9020    let gscale = |g: usize| f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
9021
9022    if a8w8_enabled() {
9023        let acts: Vec<SplitAct> = (0..b)
9024            .map(|bi| split_act(&xs_all[bi * cols..(bi + 1) * cols]))
9025            .collect();
9026        let acts = &acts;
9027        let out_addr = SendMut(out.as_mut_ptr());
9028        let run = move |start: usize, end: usize| {
9029            ROW_I8.with(|rb| {
9030                let mut buf = rb.borrow_mut();
9031                buf.resize(cols, 0);
9032                for r in start..end {
9033                    // Unpack the row's nibbles to centered i8 once
9034                    // (element 2k = low nibble, 2k+1 = high — flat order,
9035                    // same as dot_q4_row_sdot's zip).
9036                    for gi in 0..gpr {
9037                        let g = r * gpr + gi;
9038                        for (k, &bt) in packed[g * 16..(g + 1) * 16].iter().enumerate() {
9039                            buf[gi * GROUP_SIZE + k * 2] = ((bt & 0x0F) as i32 - 8) as i8 as u8;
9040                            buf[gi * GROUP_SIZE + k * 2 + 1] =
9041                                (((bt >> 4) & 0x0F) as i32 - 8) as i8 as u8;
9042                        }
9043                    }
9044                    let mut bi = 0usize;
9045                    #[cfg(target_arch = "x86_64")]
9046                    if avx2_enabled() && blocked_enabled() {
9047                        while bi + 4 <= acts.len() {
9048                            let xs = [
9049                                acts[bi].xq.as_slice(),
9050                                acts[bi + 1].xq.as_slice(),
9051                                acts[bi + 2].xq.as_slice(),
9052                                acts[bi + 3].xq.as_slice(),
9053                            ];
9054                            let d = unsafe {
9055                                if vnni_tiles_enabled() {
9056                                    dot_q4b_row_1x4_vnni(&buf, scales, r * gpr, gpr, xs)
9057                                } else {
9058                                    dot_q4b_row_1x4_avx2(&buf, scales, r * gpr, gpr, xs)
9059                                }
9060                            };
9061                            for k in 0..4 {
9062                                let act = &acts[bi + k];
9063                                let mut acc = d[k] * act.sx;
9064                                for &(j, xv) in &act.outliers {
9065                                    acc += (buf[j] as i8) as f32
9066                                        * gscale((r * cols + j) / GROUP_SIZE)
9067                                        * xv;
9068                                }
9069                                // SAFETY: disjoint (bi, r) cells per worker.
9070                                unsafe { *out_addr.at((bi + k) * rows + r) = acc };
9071                            }
9072                            bi += 4;
9073                        }
9074                    }
9075                    while bi < acts.len() {
9076                        let act = &acts[bi];
9077                        let mut acc = 0f32;
9078                        for gi in 0..gpr {
9079                            let d = dot_i8_i8(
9080                                &buf[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE],
9081                                &act.xq[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE],
9082                            );
9083                            acc += d as f32 * gscale(r * gpr + gi);
9084                        }
9085                        acc *= act.sx;
9086                        // xq is zeroed at outlier slots — exact terms.
9087                        for &(j, xv) in &act.outliers {
9088                            acc += (buf[j] as i8) as f32 * gscale((r * cols + j) / GROUP_SIZE) * xv;
9089                        }
9090                        // SAFETY: disjoint (bi, r) cells per worker row range.
9091                        unsafe { *out_addr.at(bi * rows + r) = acc };
9092                        bi += 1;
9093                    }
9094                }
9095            })
9096        };
9097        dispatch_rows(pool, rows, &run);
9098        return;
9099    }
9100
9101    let out_addr = SendMut(out.as_mut_ptr());
9102    let run = move |start: usize, end: usize| {
9103        ROW_F32.with(|rb| {
9104            let mut buf = rb.borrow_mut();
9105            buf.resize(cols, 0.0);
9106            for r in start..end {
9107                // Decode raw (nib − 8) values once; scales stay per-group
9108                // so the accumulation order matches q4matvec bit-for-bit.
9109                for gi in 0..gpr {
9110                    let g = r * gpr + gi;
9111                    for (k, &bt) in packed[g * 16..(g + 1) * 16].iter().enumerate() {
9112                        buf[gi * GROUP_SIZE + k * 2] = (bt & 0x0F) as f32 - 8.0;
9113                        buf[gi * GROUP_SIZE + k * 2 + 1] = ((bt >> 4) & 0x0F) as f32 - 8.0;
9114                    }
9115                }
9116                for bi in 0..b {
9117                    let x = &xs_all[bi * cols..(bi + 1) * cols];
9118                    let mut acc = 0f32;
9119                    for gi in 0..gpr {
9120                        let mut ga = 0f32;
9121                        // Pairwise (lo + hi) addition, matching
9122                        // q4matvec's `ga += lo·x + hi·x` shape exactly —
9123                        // a flat one-per-element loop rounds differently
9124                        // and broke bit-parity on the scalar (x86) path.
9125                        for k in 0..GROUP_SIZE / 2 {
9126                            let e = gi * GROUP_SIZE + k * 2;
9127                            ga += buf[e] * x[e] + buf[e + 1] * x[e + 1];
9128                        }
9129                        acc += ga * gscale(r * gpr + gi);
9130                    }
9131                    // SAFETY: disjoint (bi, r) cells per worker row range.
9132                    unsafe { *out_addr.at(bi * rows + r) = acc };
9133                }
9134            }
9135        })
9136    };
9137    dispatch_rows(pool, rows, &run);
9138}
9139
9140/// Batched vbit matmat: each variable-bit row is decoded from the mmap
9141/// ONCE for the whole microbatch. Same per-position math as
9142/// `vbitmatvec` (SDOT A8W8 with exact outliers / exact f32 for b=8 rows
9143/// and the scalar path).
9144#[allow(clippy::too_many_arguments)]
9145fn vbitmatmat(
9146    bytes: &[u8],
9147    offsets: &[usize],
9148    xs_all: &[f32],
9149    b: usize,
9150    rows: usize,
9151    cols: usize,
9152    out: &mut [f32],
9153    pool: Option<&Pool>,
9154) {
9155    debug_assert_eq!(xs_all.len(), b * cols);
9156    debug_assert_eq!(out.len(), b * rows);
9157    debug_assert_eq!(offsets.len(), rows + 1);
9158    let ng = cols / GROUP_SIZE;
9159    let bits = &bytes[..rows];
9160    let sc_off = rows;
9161    let gscale = |r: usize, g: usize| {
9162        let so = (r * ng + g) * 2;
9163        f16_to_f32(u16::from_le_bytes([
9164            bytes[sc_off + so],
9165            bytes[sc_off + so + 1],
9166        ]))
9167    };
9168
9169    // Decode row r's raw (u − L) values into `dst` (f32, unscaled).
9170    let decode_f32 = |r: usize, dst: &mut [f32]| {
9171        let bw = bits[r] as usize;
9172        let l = ((1i32 << (bw - 1)) - 1) as f32;
9173        let data = &bytes[offsets[r]..offsets[r + 1]];
9174        let (mut acc, mut nbits, mut idx) = (0u64, 0usize, 0usize);
9175        for d in dst.iter_mut() {
9176            while nbits < bw {
9177                acc = (acc << 8) | data[idx] as u64;
9178                idx += 1;
9179                nbits += 8;
9180            }
9181            let u = ((acc >> (nbits - bw)) & ((1u64 << bw) - 1)) as f32;
9182            nbits -= bw;
9183            *d = u - l;
9184        }
9185    };
9186
9187    if a8w8_enabled() {
9188        let acts: Vec<SplitAct> = (0..b)
9189            .map(|bi| split_act(&xs_all[bi * cols..(bi + 1) * cols]))
9190            .collect();
9191        let acts = &acts;
9192        let out_addr = SendMut(out.as_mut_ptr());
9193        let run = move |start: usize, end: usize| {
9194            for r in start..end {
9195                let bw = bits[r] as usize;
9196                if bw == 8 {
9197                    // u−L reaches 128 → no i8 path; decode once, exact
9198                    // f32 dots for every position (same as vbitmatvec).
9199                    ROW_F32.with(|rb| {
9200                        let mut buf = rb.borrow_mut();
9201                        buf.resize(cols, 0.0);
9202                        decode_f32(r, &mut buf);
9203                        for bi in 0..b {
9204                            let x = &xs_all[bi * cols..(bi + 1) * cols];
9205                            let mut dot = 0f32;
9206                            for g in 0..ng {
9207                                let mut gd = 0f32;
9208                                for k in 0..GROUP_SIZE {
9209                                    gd += buf[g * GROUP_SIZE + k] * x[g * GROUP_SIZE + k];
9210                                }
9211                                dot += gd * gscale(r, g);
9212                            }
9213                            // SAFETY: disjoint (bi, r) cells per worker range.
9214                            unsafe { *out_addr.at(bi * rows + r) = dot };
9215                        }
9216                    });
9217                    continue;
9218                }
9219                let l = (1i32 << (bw - 1)) - 1;
9220                let data = &bytes[offsets[r]..offsets[r + 1]];
9221                ROW_I8.with(|rb| {
9222                    let mut buf = rb.borrow_mut();
9223                    buf.resize(cols, 0);
9224                    #[inline(always)]
9225                    fn fill<const B: usize>(data: &[u8], l: i32, buf: &mut [u8]) {
9226                        for (blk, chunk) in buf.chunks_exact_mut(8).enumerate() {
9227                            let u = unpack8::<B>(&data[blk * B..]);
9228                            for k in 0..8 {
9229                                chunk[k] = (u[k] - l) as i8 as u8;
9230                            }
9231                        }
9232                    }
9233                    match bw {
9234                        3 => fill::<3>(data, l, &mut buf),
9235                        4 => vbit_fill4(data, &mut buf),
9236                        5 => fill::<5>(data, l, &mut buf),
9237                        6 => fill::<6>(data, l, &mut buf),
9238                        _ => unreachable!("vbit bit-width {bw} (validated at load)"),
9239                    }
9240                    let mut bi = 0usize;
9241                    // The vbit scale table shares q4_block's layout
9242                    // (contiguous f16 per (row·ng + g)), so the same
9243                    // blocked 1×4 kernel serves the decoded row.
9244                    #[cfg(target_arch = "x86_64")]
9245                    if avx2_enabled() && blocked_enabled() {
9246                        while bi + 4 <= acts.len() {
9247                            let xs = [
9248                                acts[bi].xq.as_slice(),
9249                                acts[bi + 1].xq.as_slice(),
9250                                acts[bi + 2].xq.as_slice(),
9251                                acts[bi + 3].xq.as_slice(),
9252                            ];
9253                            let sxs = [
9254                                acts[bi].sx,
9255                                acts[bi + 1].sx,
9256                                acts[bi + 2].sx,
9257                                acts[bi + 3].sx,
9258                            ];
9259                            let d = unsafe {
9260                                if vnni_tiles_enabled() {
9261                                    dot_q4b_row_1x4_sx_vnni(
9262                                        &buf,
9263                                        &bytes[sc_off..],
9264                                        r * ng,
9265                                        ng,
9266                                        xs,
9267                                        sxs,
9268                                    )
9269                                } else {
9270                                    dot_q4b_row_1x4_sx_avx2(
9271                                        &buf,
9272                                        &bytes[sc_off..],
9273                                        r * ng,
9274                                        ng,
9275                                        xs,
9276                                        sxs,
9277                                    )
9278                                }
9279                            };
9280                            for k in 0..4 {
9281                                let act = &acts[bi + k];
9282                                let mut dot = d[k];
9283                                for &(j, xv) in &act.outliers {
9284                                    dot += (buf[j] as i8) as f32 * gscale(r, j / GROUP_SIZE) * xv;
9285                                }
9286                                // SAFETY: disjoint (bi, r) cells per worker.
9287                                unsafe { *out_addr.at((bi + k) * rows + r) = dot };
9288                            }
9289                            bi += 4;
9290                        }
9291                    }
9292                    while bi < acts.len() {
9293                        let act = &acts[bi];
9294                        let mut dot = 0f32;
9295                        for g in 0..ng {
9296                            let d = dot_i8_i8(
9297                                &buf[g * GROUP_SIZE..(g + 1) * GROUP_SIZE],
9298                                &act.xq[g * GROUP_SIZE..(g + 1) * GROUP_SIZE],
9299                            ) as f32
9300                                * act.sx;
9301                            dot += d * gscale(r, g);
9302                        }
9303                        for &(j, xv) in &act.outliers {
9304                            dot += (buf[j] as i8) as f32 * gscale(r, j / GROUP_SIZE) * xv;
9305                        }
9306                        // SAFETY: disjoint (bi, r) cells per worker range.
9307                        unsafe { *out_addr.at(bi * rows + r) = dot };
9308                        bi += 1;
9309                    }
9310                });
9311            }
9312        };
9313        dispatch_rows(pool, rows, &run);
9314        return;
9315    }
9316
9317    let out_addr = SendMut(out.as_mut_ptr());
9318    let run = move |start: usize, end: usize| {
9319        ROW_F32.with(|rb| {
9320            let mut buf = rb.borrow_mut();
9321            buf.resize(cols, 0.0);
9322            for r in start..end {
9323                decode_f32(r, &mut buf);
9324                for bi in 0..b {
9325                    let x = &xs_all[bi * cols..(bi + 1) * cols];
9326                    let mut dot = 0f32;
9327                    for g in 0..ng {
9328                        let mut gd = 0f32;
9329                        for k in 0..GROUP_SIZE {
9330                            gd += buf[g * GROUP_SIZE + k] * x[g * GROUP_SIZE + k];
9331                        }
9332                        dot += gd * gscale(r, g);
9333                    }
9334                    // SAFETY: disjoint (bi, r) cells per worker range.
9335                    unsafe { *out_addr.at(bi * rows + r) = dot };
9336                }
9337            }
9338        })
9339    };
9340    dispatch_rows(pool, rows, &run);
9341}
9342
9343/// Build a GPU batch job for a q8-family mapped tensor (primary
9344/// shard): prescaled input + directory coordinates. None → not
9345/// GPU-eligible, caller stays on the CPU.
9346pub(crate) fn gpu_batch_job<'a>(
9347    t: &'a QTensor,
9348    x: &[f32],
9349) -> Option<(std::sync::Arc<CmfModel>, crate::gpu::BatchJob<'a>)> {
9350    match t {
9351        QTensor::Mapped {
9352            model,
9353            idx,
9354            dtype: dt @ (TensorDtype::Q8Row | TensorDtype::Q8_2f),
9355            rows,
9356            cols,
9357            row_scale,
9358            col_field,
9359            ..
9360        } => Some((
9361            model.clone(),
9362            crate::gpu::BatchJob {
9363                idx: *idx,
9364                rows: *rows,
9365                cols: *cols,
9366                row_scale,
9367                xs: prescale(x, col_field, *dt).into_owned(),
9368                layout: crate::gpu::BatchLayout::Q8,
9369            },
9370        )),
9371        // q1: raw f32 activations, tile-embedded scales.
9372        QTensor::Mapped {
9373            model,
9374            idx,
9375            dtype: TensorDtype::Q1,
9376            rows,
9377            cols,
9378            ..
9379        } => Some((
9380            model.clone(),
9381            crate::gpu::BatchJob {
9382                idx: *idx,
9383                rows: *rows,
9384                cols: *cols,
9385                row_scale: &[],
9386                xs: x.to_vec(),
9387                layout: crate::gpu::BatchLayout::Q1,
9388            },
9389        )),
9390        // q4_tiled / q4tp: raw f32 activations; the scales live in the
9391        // payload (inline tiles / row ladder), so row_scale stays empty.
9392        // The GDN projection batch already runs these layouts on Metal —
9393        // this arm lets the attention QKV batch reach the same kernels.
9394        QTensor::Mapped {
9395            model,
9396            idx,
9397            dtype: dt @ (TensorDtype::Q4Tiled | TensorDtype::Q4TiledP),
9398            rows,
9399            cols,
9400            ..
9401        } => Some((
9402            model.clone(),
9403            crate::gpu::BatchJob {
9404                idx: *idx,
9405                rows: *rows,
9406                cols: *cols,
9407                row_scale: &[],
9408                xs: x.to_vec(),
9409                layout: if *dt == TensorDtype::Q4Tiled {
9410                    crate::gpu::BatchLayout::Q4t
9411                } else {
9412                    crate::gpu::BatchLayout::Q4tp
9413                },
9414            },
9415        )),
9416        _ => None,
9417    }
9418}
9419
9420thread_local! {
9421    static PRESCALE_BUF1: std::cell::RefCell<Vec<f32>> = const { std::cell::RefCell::new(Vec::new()) };
9422    static PRESCALE_BUF2: std::cell::RefCell<Vec<f32>> = const { std::cell::RefCell::new(Vec::new()) };
9423}
9424
9425pub(crate) fn prescale<'a>(
9426    x: &'a [f32],
9427    col_field: &[f32],
9428    dtype: TensorDtype,
9429) -> std::borrow::Cow<'a, [f32]> {
9430    if dtype == TensorDtype::Q8_2f {
9431        x.iter().zip(col_field).map(|(a, c)| a * c).collect()
9432    } else {
9433        std::borrow::Cow::Borrowed(x)
9434    }
9435}
9436
9437/// θ col-field fold for q8_2f activations. Borrowed pass-through for
9438/// every other dtype, using thread-local buffers to eliminate per-matvec allocations.
9439pub(crate) fn prescale_with<R, F: FnOnce(&[f32]) -> R>(
9440    x: &[f32],
9441    col_field: &[f32],
9442    dtype: TensorDtype,
9443    buf_id: u8,
9444    f: F,
9445) -> R {
9446    if dtype == TensorDtype::Q8_2f {
9447        if buf_id == 1 {
9448            PRESCALE_BUF1.with(|b| {
9449                let mut buf = b.borrow_mut();
9450                buf.clear();
9451                buf.extend(x.iter().zip(col_field).map(|(a, c)| a * c));
9452                f(&buf)
9453            })
9454        } else {
9455            PRESCALE_BUF2.with(|b| {
9456                let mut buf = b.borrow_mut();
9457                buf.clear();
9458                buf.extend(x.iter().zip(col_field).map(|(a, c)| a * c));
9459                f(&buf)
9460            })
9461        }
9462    } else {
9463        f(x)
9464    }
9465}
9466
9467// ───────────────────── x86-64 AVX2 kernels (roadmap этап 2) ─────────────────────
9468
9469/// AVX2+FMA available? Default ON when the CPU supports both;
9470/// `CMF_AVX2=0` disables (falls back to the autovectorized loops).
9471#[cfg(target_arch = "x86_64")]
9472pub(crate) fn avx2_enabled() -> bool {
9473    use std::sync::OnceLock;
9474    static ON: OnceLock<bool> = OnceLock::new();
9475    *ON.get_or_init(|| {
9476        std::env::var("CMF_AVX2").map(|v| v != "0").unwrap_or(true)
9477            && std::arch::is_x86_feature_detected!("avx2")
9478            && std::arch::is_x86_feature_detected!("fma")
9479    })
9480}
9481
9482/// AVX2 A8W8 allowed? The quantized-activation contract is switched by
9483/// the SAME env as the ARM SDOT path: `CMF_SDOT=0` keeps exact kernels
9484/// (the golden-parity exact gate relies on it) — AVX2 f32 kernels stay
9485/// active either way, they are exact (regrouped sums only).
9486#[cfg(target_arch = "x86_64")]
9487fn avx2_a8w8_enabled() -> bool {
9488    if FLOAT_ACTIVATIONS.get() {
9489        return false;
9490    }
9491    use std::sync::OnceLock;
9492    static ON: OnceLock<bool> = OnceLock::new();
9493    *ON.get_or_init(|| {
9494        avx2_enabled() && std::env::var("CMF_SDOT").map(|v| v != "0").unwrap_or(true)
9495    })
9496}
9497
9498thread_local! {
9499    static FULL_GPU_Q8: std::cell::Cell<bool> = const { std::cell::Cell::new(false) };
9500}
9501
9502/// Match the graph's full-device q8 projection precision on MiMo's host
9503/// tail. Does not enable the GPU or bypass a CPU-only/device-refusal gate.
9504pub(crate) fn enter_full_gpu_q8_scope() -> impl Drop {
9505    struct Restore(bool, std::marker::PhantomData<std::rc::Rc<()>>);
9506    impl Drop for Restore {
9507        fn drop(&mut self) {
9508            FULL_GPU_Q8.set(self.0);
9509        }
9510    }
9511    Restore(FULL_GPU_Q8.replace(true), std::marker::PhantomData)
9512}
9513
9514// Dynamic MiMo experts must not change activation precision when a cache
9515// fill moves them from CPU to GPU. Thread-local: only the cold-expert
9516// dispatch selects float kernels; concurrent pipelines keep their policy.
9517thread_local! {
9518    static FLOAT_ACTIVATIONS: std::cell::Cell<bool> = const { std::cell::Cell::new(false) };
9519}
9520
9521pub(crate) fn float_activations_scope<R>(f: impl FnOnce() -> R) -> R {
9522    struct Restore(bool);
9523    impl Drop for Restore {
9524        fn drop(&mut self) {
9525            FLOAT_ACTIVATIONS.set(self.0);
9526        }
9527    }
9528    let _restore = Restore(FLOAT_ACTIVATIONS.replace(true));
9529    f()
9530}
9531
9532/// Row-exact batching: while set, the x86 batched kernels (`qmatmat`,
9533/// `q4tp_matmat`) compute every (weight row, token) cell with the
9534/// single-token kernel instead of the blocked 2×4 / 1×4 / 1×8 tiles, so a
9535/// token's result does not depend on the batch it rides in and equals its
9536/// matvec. The MiMo speculative verify holds it (`row_exact_scope`) — its
9537/// accepted rows must be the rows plain decode would have produced.
9538// Shared with pool workers, so overlapping requests must keep the mode
9539// enabled until the LAST scope leaves. Saving/restoring a global bool is
9540// incorrect when two threads enter and leave in a non-LIFO order.
9541static ROW_EXACT: std::sync::atomic::AtomicUsize = std::sync::atomic::AtomicUsize::new(0);
9542
9543pub(crate) fn row_exact() -> bool {
9544    ROW_EXACT.load(std::sync::atomic::Ordering::Acquire) != 0
9545}
9546
9547fn counted_row_exact_scope<R>(active: &std::sync::atomic::AtomicUsize, f: impl FnOnce() -> R) -> R {
9548    struct Restore<'a>(&'a std::sync::atomic::AtomicUsize);
9549    impl Drop for Restore<'_> {
9550        fn drop(&mut self) {
9551            self.0.fetch_sub(1, std::sync::atomic::Ordering::AcqRel);
9552        }
9553    }
9554    active.fetch_add(1, std::sync::atomic::Ordering::AcqRel);
9555    let _restore = Restore(active);
9556    f()
9557}
9558
9559/// Run `f` with row-exact batching on (also released on unwind).
9560pub(crate) fn row_exact_scope<R>(f: impl FnOnce() -> R) -> R {
9561    counted_row_exact_scope(&ROW_EXACT, f)
9562}
9563
9564/// A8W8 quantized-activation path available on THIS machine? One
9565/// switch across architectures: ARM dotprod (CMF_SDOT) or x86 AVX2
9566/// (CMF_AVX2 + the same CMF_SDOT exact-contract override).
9567#[inline]
9568pub(crate) fn a8w8_enabled() -> bool {
9569    #[cfg(target_arch = "aarch64")]
9570    {
9571        sdot_enabled()
9572    }
9573    #[cfg(target_arch = "x86_64")]
9574    {
9575        avx2_a8w8_enabled()
9576    }
9577    #[cfg(not(any(target_arch = "aarch64", target_arch = "x86_64")))]
9578    {
9579        false
9580    }
9581}
9582
9583/// int8·int8 dot dispatch: SDOT on ARM; AVX-512 VNNI (vpdpbusd) or AVX2
9584/// maddubs on x86. Callers are gated by `a8w8_enabled()`.
9585#[inline]
9586#[allow(unreachable_code)]
9587fn dot_i8_i8(w: &[u8], xq: &[i8]) -> i32 {
9588    #[cfg(target_arch = "aarch64")]
9589    unsafe {
9590        return dot_i8_sdot(w, xq);
9591    }
9592    #[cfg(target_arch = "x86_64")]
9593    unsafe {
9594        if avx512vnni_enabled() {
9595            return dot_i8_i8_vnni(w, xq);
9596        }
9597        return dot_i8_i8_avx2(w, xq);
9598    }
9599    w.iter()
9600        .zip(xq)
9601        .map(|(&a, &b)| (a as i8) as i32 * b as i32)
9602        .sum()
9603}
9604
9605/// AVX-512 VNNI available? (F+BW+VL+VNNI; `CMF_AVX512=0` falls back to
9606/// AVX2.) VL matters: short 32-byte groups (q4/vbit) ride the 256-bit
9607/// `vpdpbusd` encoding.
9608#[cfg(target_arch = "x86_64")]
9609fn avx512vnni_enabled() -> bool {
9610    use std::sync::OnceLock;
9611    static ON: OnceLock<bool> = OnceLock::new();
9612    *ON.get_or_init(|| {
9613        std::env::var("CMF_AVX512")
9614            .map(|v| v != "0")
9615            .unwrap_or(true)
9616            && std::arch::is_x86_feature_detected!("avx512f")
9617            && std::arch::is_x86_feature_detected!("avx512bw")
9618            && std::arch::is_x86_feature_detected!("avx512vl")
9619            && std::arch::is_x86_feature_detected!("avx512vnni")
9620    })
9621}
9622
9623/// Grouped-codec VNNI arms (the q4t/q4b/q1/q1t tile kernels): default
9624/// ON where AVX-512 VNNI exists (`CMF_VNNI_TILES=0` opt-out). Measured
9625/// on Ryzen 7950X (Zen4, 3 alternating process pairs, blocked GEMM
9626/// 4864×896 b=256): q4t 63→68 GF/s (+8%), q1 53→56 (+6%), q4b 72→75
9627/// (+4%) — consistent, no leg regressed. The tile kernels keep a
9628/// horizontal reduce per 32-weight group, so the `vpdpbusd` saving is
9629/// smaller than the long-dot q8 win (+13%), but it is real and free.
9630#[cfg(target_arch = "x86_64")]
9631fn vnni_tiles_enabled() -> bool {
9632    use std::sync::OnceLock;
9633    static ON: OnceLock<bool> = OnceLock::new();
9634    *ON.get_or_init(|| {
9635        std::env::var("CMF_VNNI_TILES")
9636            .map(|v| v != "0")
9637            .unwrap_or(true)
9638            && avx512vnni_enabled()
9639    })
9640}
9641
9642/// One 256-bit u8×i8 dot → i32 via `vpdpbusd` into a fresh accumulator
9643/// plus the same horizontal reduce the AVX2 kernels use. Products are
9644/// bounded (|w| ≤ 8 or ≤ 1), so maddubs never saturated — the i32 sum
9645/// is bit-identical to the maddubs+madd pair it replaces.
9646#[cfg(target_arch = "x86_64")]
9647#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
9648#[inline]
9649unsafe fn dpbusd_hsum(aw: core::arch::x86_64::__m256i, xs: core::arch::x86_64::__m256i) -> i32 {
9650    // SAFETY: pure register math.
9651    unsafe {
9652        use core::arch::x86_64::*;
9653        let d = _mm256_dpbusd_epi32(_mm256_setzero_si256(), aw, xs);
9654        let hi128 = _mm256_extracti128_si256::<1>(d);
9655        let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
9656        let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
9657        let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
9658        _mm_cvtsi128_si32(s32)
9659    }
9660}
9661
9662/// int8·int8 via AVX-512 VNNI: `vpdpbusd` fuses the maddubs+madd+add
9663/// triple into one u8×i8 dot-accumulate. AVX-512 has no vpsignb, so the
9664/// |w|·sign(x,w) trick becomes |w| × (x negated where w<0) via a mask
9665/// subtract — w==0 lanes contribute 0 through |w|=0 either way.
9666#[cfg(target_arch = "x86_64")]
9667#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
9668unsafe fn dot_i8_i8_vnni(w: &[u8], xq: &[i8]) -> i32 {
9669    // SAFETY: callers uphold slice-length contracts (see call sites).
9670    unsafe {
9671        use core::arch::x86_64::*;
9672        let n = w.len();
9673        let mut j = 0usize;
9674        let mut total: i32;
9675        // 4 independent accumulators: vpdpbusd is its own loop-carried
9676        // dependency (~5-cycle latency) — a single-acc loop runs
9677        // latency-bound and LOSES to the AVX2 maddubs kernel, measured
9678        // on Granite Rapids.
9679        {
9680            #[inline(always)]
9681            unsafe fn step(
9682                w: *const u8,
9683                x: *const i8,
9684                acc: core::arch::x86_64::__m512i,
9685            ) -> core::arch::x86_64::__m512i {
9686                unsafe {
9687                    use core::arch::x86_64::*;
9688                    let wv = _mm512_loadu_si512(w as *const _);
9689                    let xv = _mm512_loadu_si512(x as *const _);
9690                    let aw = _mm512_abs_epi8(wv);
9691                    let neg = _mm512_movepi8_mask(wv);
9692                    let sx = _mm512_mask_sub_epi8(xv, neg, _mm512_setzero_si512(), xv);
9693                    _mm512_dpbusd_epi32(acc, aw, sx)
9694                }
9695            }
9696            let (mut a0, mut a1, mut a2, mut a3) = (
9697                _mm512_setzero_si512(),
9698                _mm512_setzero_si512(),
9699                _mm512_setzero_si512(),
9700                _mm512_setzero_si512(),
9701            );
9702            while j + 256 <= n {
9703                a0 = step(w.as_ptr().add(j), xq.as_ptr().add(j), a0);
9704                a1 = step(w.as_ptr().add(j + 64), xq.as_ptr().add(j + 64), a1);
9705                a2 = step(w.as_ptr().add(j + 128), xq.as_ptr().add(j + 128), a2);
9706                a3 = step(w.as_ptr().add(j + 192), xq.as_ptr().add(j + 192), a3);
9707                j += 256;
9708            }
9709            while j + 64 <= n {
9710                a0 = step(w.as_ptr().add(j), xq.as_ptr().add(j), a0);
9711                j += 64;
9712            }
9713            let s01 = _mm512_add_epi32(a0, a1);
9714            let s23 = _mm512_add_epi32(a2, a3);
9715            total = _mm512_reduce_add_epi32(_mm512_add_epi32(s01, s23));
9716        }
9717        // 32-wide (q4/vbit groups are exactly 32 bytes).
9718        if j + 32 <= n {
9719            let wv = _mm256_loadu_si256(w.as_ptr().add(j) as *const __m256i);
9720            let xv = _mm256_loadu_si256(xq.as_ptr().add(j) as *const __m256i);
9721            let d = _mm256_dpbusd_epi32(
9722                _mm256_setzero_si256(),
9723                _mm256_abs_epi8(wv),
9724                _mm256_sign_epi8(xv, wv),
9725            );
9726            let hi128 = _mm256_extracti128_si256::<1>(d);
9727            let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
9728            let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
9729            let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
9730            total += _mm_cvtsi128_si32(s32);
9731            j += 32;
9732        }
9733        while j < n {
9734            total += (w[j] as i8) as i32 * xq[j] as i32;
9735            j += 1;
9736        }
9737        total
9738    }
9739}
9740
9741/// i8 row · f32 x via AVX2/FMA (x86 mirror of `dot_i8_f32_neon`).
9742#[cfg(target_arch = "x86_64")]
9743#[target_feature(enable = "avx2,fma")]
9744unsafe fn dot_i8_f32_avx2(w: &[u8], x: &[f32]) -> f32 {
9745    // SAFETY: callers uphold slice-length contracts (see call sites).
9746    unsafe {
9747        use core::arch::x86_64::*;
9748        let n = x.len();
9749        let wp = w.as_ptr();
9750        let xp = x.as_ptr();
9751        let (mut a0, mut a1) = (_mm256_setzero_ps(), _mm256_setzero_ps());
9752        let mut j = 0usize;
9753        while j + 16 <= n {
9754            let wb = _mm_loadu_si128(wp.add(j) as *const __m128i);
9755            let lo = _mm256_cvtepi8_epi32(wb);
9756            let hi = _mm256_cvtepi8_epi32(_mm_srli_si128::<8>(wb));
9757            a0 = _mm256_fmadd_ps(_mm256_cvtepi32_ps(lo), _mm256_loadu_ps(xp.add(j)), a0);
9758            a1 = _mm256_fmadd_ps(_mm256_cvtepi32_ps(hi), _mm256_loadu_ps(xp.add(j + 8)), a1);
9759            j += 16;
9760        }
9761        let acc = _mm256_add_ps(a0, a1);
9762        let hi128 = _mm256_extractf128_ps::<1>(acc);
9763        let s128 = _mm_add_ps(_mm256_castps256_ps128(acc), hi128);
9764        let s64 = _mm_add_ps(s128, _mm_movehl_ps(s128, s128));
9765        let s32 = _mm_add_ss(s64, _mm_shuffle_ps::<1>(s64, s64));
9766        let mut sum = _mm_cvtss_f32(s32);
9767        while j < n {
9768            sum += (*wp.add(j) as i8) as f32 * *xp.add(j);
9769            j += 1;
9770        }
9771        sum
9772    }
9773}
9774
9775/// int8(weight)·int8(activation) → i32 via AVX2 maddubs — the x86
9776/// analogue of the SDOT A8W8 path. `maddubs` takes u8×i8, so the
9777/// standard sign trick applies: |w| × sign(x, w) ≡ w × x per lane.
9778/// Pair saturation is safe: |w|≤128, |x|≤127 → 2·128·127 < 32767.
9779#[cfg(target_arch = "x86_64")]
9780#[target_feature(enable = "avx2")]
9781unsafe fn dot_i8_i8_avx2(w: &[u8], xq: &[i8]) -> i32 {
9782    // SAFETY: callers uphold slice-length contracts (see call sites).
9783    unsafe {
9784        use core::arch::x86_64::*;
9785        let n = w.len();
9786        let ones = _mm256_set1_epi16(1);
9787        let mut acc = _mm256_setzero_si256();
9788        let mut j = 0usize;
9789        while j + 32 <= n {
9790            let wv = _mm256_loadu_si256(w.as_ptr().add(j) as *const __m256i);
9791            let xv = _mm256_loadu_si256(xq.as_ptr().add(j) as *const __m256i);
9792            let p16 = _mm256_maddubs_epi16(_mm256_abs_epi8(wv), _mm256_sign_epi8(xv, wv));
9793            acc = _mm256_add_epi32(acc, _mm256_madd_epi16(p16, ones));
9794            j += 32;
9795        }
9796        let hi128 = _mm256_extracti128_si256::<1>(acc);
9797        let s128 = _mm_add_epi32(_mm256_castsi256_si128(acc), hi128);
9798        let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
9799        let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
9800        let mut s = _mm_cvtsi128_si32(s32);
9801        while j < n {
9802            s += (w[j] as i8) as i32 * xq[j] as i32;
9803            j += 1;
9804        }
9805        s
9806    }
9807}
9808
9809/// smmla 2×4: one instruction covers a 2-row × 2-activation × 8-deep
9810/// tile (32 MACs vs sdot's 16) — the weight pair loads once per 8-k
9811/// slice as a combined 2×8 register and meets two activation pairs.
9812#[cfg(target_arch = "aarch64")]
9813#[target_feature(enable = "neon,i8mm")]
9814unsafe fn dot_i8_smmla_2x4(w0: &[u8], w1: &[u8], xs: [&[i8]; 4]) -> [[i32; 4]; 2] {
9815    // SAFETY: callers uphold slice-length contracts.
9816    unsafe {
9817        use core::arch::aarch64::*;
9818        use core::arch::asm;
9819        let n = w0.len();
9820        let w0p = w0.as_ptr() as *const i8;
9821        let w1p = w1.as_ptr() as *const i8;
9822        // acc01 holds [c(r0,x0) c(r0,x1) c(r1,x0) c(r1,x1)]; acc23 the
9823        // same for x2/x3.
9824        let mut acc01 = vdupq_n_s32(0);
9825        let mut acc23 = vdupq_n_s32(0);
9826        let mut i = 0usize;
9827        while i + 8 <= n {
9828            let wa = vcombine_s8(vld1_s8(w0p.add(i)), vld1_s8(w1p.add(i)));
9829            let xb01 = vcombine_s8(
9830                vld1_s8(xs[0].as_ptr().add(i)),
9831                vld1_s8(xs[1].as_ptr().add(i)),
9832            );
9833            let xb23 = vcombine_s8(
9834                vld1_s8(xs[2].as_ptr().add(i)),
9835                vld1_s8(xs[3].as_ptr().add(i)),
9836            );
9837            asm!(
9838                "smmla {a01:v}.4s, {w:v}.16b, {x01:v}.16b",
9839                "smmla {a23:v}.4s, {w:v}.16b, {x23:v}.16b",
9840                a01 = inout(vreg) acc01, a23 = inout(vreg) acc23,
9841                w = in(vreg) wa, x01 = in(vreg) xb01, x23 = in(vreg) xb23,
9842                options(pure, nomem, nostack),
9843            );
9844            i += 8;
9845        }
9846        let mut out = [[0i32; 4]; 2];
9847        let a01: [i32; 4] = core::mem::transmute(acc01);
9848        let a23: [i32; 4] = core::mem::transmute(acc23);
9849        out[0][0] = a01[0];
9850        out[0][1] = a01[1];
9851        out[1][0] = a01[2];
9852        out[1][1] = a01[3];
9853        out[0][2] = a23[0];
9854        out[0][3] = a23[1];
9855        out[1][2] = a23[2];
9856        out[1][3] = a23[3];
9857        if i < n {
9858            for (k, x) in xs.iter().enumerate() {
9859                for j in i..n {
9860                    out[0][k] += (w0[j] as i8) as i32 * x[j] as i32;
9861                    out[1][k] += (w1[j] as i8) as i32 * x[j] as i32;
9862                }
9863            }
9864        }
9865        out
9866    }
9867}
9868
9869/// ARM twin of the x86 blocked prefill GEMM: two weight rows stay in
9870/// registers across four activation streams, eight sdot accumulators.
9871/// (The per-row form re-read each W row once per activation.)
9872#[cfg(target_arch = "aarch64")]
9873#[target_feature(enable = "neon,dotprod")]
9874unsafe fn dot_i8_sdot_2x4(w0: &[u8], w1: &[u8], xs: [&[i8]; 4]) -> [[i32; 4]; 2] {
9875    // SAFETY: callers uphold slice-length contracts.
9876    unsafe {
9877        use core::arch::aarch64::*;
9878        use core::arch::asm;
9879        let n = w0.len();
9880        let w0p = w0.as_ptr() as *const i8;
9881        let w1p = w1.as_ptr() as *const i8;
9882        let mut acc = [[vdupq_n_s32(0); 4]; 2];
9883        let mut i = 0usize;
9884        while i + 16 <= n {
9885            let wv0 = vld1q_s8(w0p.add(i));
9886            let wv1 = vld1q_s8(w1p.add(i));
9887            for (k, x) in xs.iter().enumerate() {
9888                let xv = vld1q_s8(x.as_ptr().add(i));
9889                let (mut a0, mut a1) = (acc[0][k], acc[1][k]);
9890                asm!(
9891                    "sdot {a0:v}.4s, {w0:v}.16b, {x:v}.16b",
9892                    "sdot {a1:v}.4s, {w1:v}.16b, {x:v}.16b",
9893                    a0 = inout(vreg) a0, a1 = inout(vreg) a1,
9894                    w0 = in(vreg) wv0, w1 = in(vreg) wv1, x = in(vreg) xv,
9895                    options(pure, nomem, nostack),
9896                );
9897                acc[0][k] = a0;
9898                acc[1][k] = a1;
9899            }
9900            i += 16;
9901        }
9902        let mut out = [[0i32; 4]; 2];
9903        for r in 0..2 {
9904            for k in 0..4 {
9905                out[r][k] = vaddvq_s32(acc[r][k]);
9906            }
9907        }
9908        if i < n {
9909            for (k, x) in xs.iter().enumerate() {
9910                for j in i..n {
9911                    out[0][k] += (w0[j] as i8) as i32 * x[j] as i32;
9912                    out[1][k] += (w1[j] as i8) as i32 * x[j] as i32;
9913                }
9914            }
9915        }
9916        out
9917    }
9918}
9919
9920/// Blocked 2 weight rows × 4 activations for the prefill GEMM
9921/// (roadmap P0: packed panels + multi-row accumulators). The two rows'
9922/// abs() live in registers across all four activation streams; the
9923/// sign-fixup is recomputed per pair (the price of the maddubs trick).
9924/// Returns raw i8·i8 dots; the caller applies scales and outliers.
9925#[cfg(target_arch = "x86_64")]
9926#[target_feature(enable = "avx2")]
9927unsafe fn dot_i8_i8_avx2_2x4(w0: &[u8], w1: &[u8], xs: [&[i8]; 4]) -> [[i32; 4]; 2] {
9928    // SAFETY: callers uphold slice-length contracts.
9929    unsafe {
9930        use core::arch::x86_64::*;
9931        let n = w0.len();
9932        let ones = _mm256_set1_epi16(1);
9933        let mut acc = [[_mm256_setzero_si256(); 4]; 2];
9934        let mut j = 0usize;
9935        while j + 32 <= n {
9936            let wv0 = _mm256_loadu_si256(w0.as_ptr().add(j) as *const __m256i);
9937            let wv1 = _mm256_loadu_si256(w1.as_ptr().add(j) as *const __m256i);
9938            let aw0 = _mm256_abs_epi8(wv0);
9939            let aw1 = _mm256_abs_epi8(wv1);
9940            for (k, x) in xs.iter().enumerate() {
9941                let xv = _mm256_loadu_si256(x.as_ptr().add(j) as *const __m256i);
9942                let p0 = _mm256_maddubs_epi16(aw0, _mm256_sign_epi8(xv, wv0));
9943                acc[0][k] = _mm256_add_epi32(acc[0][k], _mm256_madd_epi16(p0, ones));
9944                let p1 = _mm256_maddubs_epi16(aw1, _mm256_sign_epi8(xv, wv1));
9945                acc[1][k] = _mm256_add_epi32(acc[1][k], _mm256_madd_epi16(p1, ones));
9946            }
9947            j += 32;
9948        }
9949        let mut out = [[0i32; 4]; 2];
9950        for r in 0..2 {
9951            for k in 0..4 {
9952                let a = acc[r][k];
9953                let hi128 = _mm256_extracti128_si256::<1>(a);
9954                let s128 = _mm_add_epi32(_mm256_castsi256_si128(a), hi128);
9955                let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
9956                let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
9957                out[r][k] = _mm_cvtsi128_si32(s32);
9958            }
9959        }
9960        if j < n {
9961            for (k, x) in xs.iter().enumerate() {
9962                for i in j..n {
9963                    out[0][k] += (w0[i] as i8) as i32 * x[i] as i32;
9964                    out[1][k] += (w1[i] as i8) as i32 * x[i] as i32;
9965                }
9966            }
9967        }
9968        out
9969    }
9970}
9971
9972/// AVX2/VNNI q8 row dot with exact outlier correction (x86 mirror of
9973/// `row_dot_sdot` — same A8W8 contract). With AVX-512 VNNI the row goes
9974/// through the bias trick: Σ(w+128)·x via pure `vpdpbusd` (no per-lane
9975/// sign fixups), corrected by −128·Σx with Σx precomputed per split.
9976#[cfg(target_arch = "x86_64")]
9977#[inline]
9978fn row_dot_avx2(row: &[u8], act: &SplitAct) -> f32 {
9979    let dot = if avx512vnni_enabled() && row.len() >= 64 {
9980        (unsafe { dot_u8p128_i8_vnni(row, &act.xq) }) - 128 * act.xsum
9981    } else {
9982        unsafe { dot_i8_i8_avx2(row, &act.xq) }
9983    };
9984    let mut acc = dot as f32 * act.sx;
9985    for &(j, xv) in &act.outliers {
9986        acc += (row[j] as i8) as f32 * xv;
9987    }
9988    acc
9989}
9990
9991/// Σ (w[i]+128)·x[i] via pure `vpdpbusd` — the caller subtracts
9992/// 128·Σx. Four independent accumulators (dpbusd is ~5-cycle latency;
9993/// a single-acc loop runs latency-bound, measured on Granite Rapids).
9994#[cfg(target_arch = "x86_64")]
9995#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
9996unsafe fn dot_u8p128_i8_vnni(w: &[u8], xq: &[i8]) -> i32 {
9997    // SAFETY: callers uphold slice-length contracts (see call sites).
9998    unsafe {
9999        use core::arch::x86_64::*;
10000        let n = w.len();
10001        let flip = _mm512_set1_epi8(-128); // XOR 0x80: i8 w → u8 (w+128)
10002        #[inline(always)]
10003        unsafe fn step(
10004            w: *const u8,
10005            x: *const i8,
10006            flip: core::arch::x86_64::__m512i,
10007            acc: core::arch::x86_64::__m512i,
10008        ) -> core::arch::x86_64::__m512i {
10009            unsafe {
10010                use core::arch::x86_64::*;
10011                let wv = _mm512_xor_si512(_mm512_loadu_si512(w as *const _), flip);
10012                _mm512_dpbusd_epi32(acc, wv, _mm512_loadu_si512(x as *const _))
10013            }
10014        }
10015        let (mut a0, mut a1, mut a2, mut a3) = (
10016            _mm512_setzero_si512(),
10017            _mm512_setzero_si512(),
10018            _mm512_setzero_si512(),
10019            _mm512_setzero_si512(),
10020        );
10021        let mut j = 0usize;
10022        while j + 256 <= n {
10023            a0 = step(w.as_ptr().add(j), xq.as_ptr().add(j), flip, a0);
10024            a1 = step(w.as_ptr().add(j + 64), xq.as_ptr().add(j + 64), flip, a1);
10025            a2 = step(w.as_ptr().add(j + 128), xq.as_ptr().add(j + 128), flip, a2);
10026            a3 = step(w.as_ptr().add(j + 192), xq.as_ptr().add(j + 192), flip, a3);
10027            j += 256;
10028        }
10029        while j + 64 <= n {
10030            a0 = step(w.as_ptr().add(j), xq.as_ptr().add(j), flip, a0);
10031            j += 64;
10032        }
10033        let mut total = _mm512_reduce_add_epi32(_mm512_add_epi32(
10034            _mm512_add_epi32(a0, a1),
10035            _mm512_add_epi32(a2, a3),
10036        ));
10037        // Scalar tail: (w as i8) + 128 ≡ (w as u8) ^ 0x80.
10038        while j < n {
10039            total += ((w[j] ^ 0x80) as i32) * xq[j] as i32;
10040            j += 1;
10041        }
10042        total
10043    }
10044}
10045
10046/// One q4 row via AVX2: nibbles → centered i8 (unpacklo/hi restores the
10047/// writer's flat order, same as the NEON vzip pair), maddubs against
10048/// the pre-quantized activation group, × the group's f16 scale. Pair
10049/// saturation safe: |w|≤8, |x|≤127 → 2·8·127 ≪ 32767. Mirror of
10050/// `dot_q4_row_sdot`.
10051#[cfg(target_arch = "x86_64")]
10052#[target_feature(enable = "avx2")]
10053unsafe fn dot_q4_row_avx2(packed: &[u8], scales: &[u8], g0: usize, gpr: usize, xq: &[i8]) -> f32 {
10054    // SAFETY: callers uphold slice-length contracts (16 packed bytes and
10055    // 2 scale bytes per group; xq.len() == gpr·GROUP_SIZE).
10056    unsafe {
10057        use core::arch::x86_64::*;
10058        let lomask = _mm_set1_epi8(0x0F);
10059        let eight = _mm256_set1_epi8(8);
10060        let ones = _mm256_set1_epi16(1);
10061        let mut acc = 0f32;
10062        for gi in 0..gpr {
10063            let g = g0 + gi;
10064            let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
10065            let b = _mm_loadu_si128(packed.as_ptr().add(g * 16) as *const __m128i);
10066            let lo = _mm_and_si128(b, lomask);
10067            let hi = _mm_and_si128(_mm_srli_epi16::<4>(b), lomask);
10068            let w = _mm256_sub_epi8(
10069                _mm256_set_m128i(_mm_unpackhi_epi8(lo, hi), _mm_unpacklo_epi8(lo, hi)),
10070                eight,
10071            );
10072            let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
10073            let p16 = _mm256_maddubs_epi16(_mm256_abs_epi8(w), _mm256_sign_epi8(x, w));
10074            let d = _mm256_madd_epi16(p16, ones);
10075            let hi128 = _mm256_extracti128_si256::<1>(d);
10076            let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
10077            let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
10078            let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
10079            acc += _mm_cvtsi128_si32(s32) as f32 * s;
10080        }
10081        acc
10082    }
10083}
10084
10085/// Two-activation q4 row via AVX2: nibbles unpacked ONCE per group,
10086/// both activations dotted against the same centered i8 register.
10087#[cfg(target_arch = "x86_64")]
10088#[target_feature(enable = "avx2")]
10089unsafe fn dot_q4_row_avx2_2(
10090    packed: &[u8],
10091    scales: &[u8],
10092    g0: usize,
10093    gpr: usize,
10094    xq1: &[i8],
10095    xq2: &[i8],
10096) -> (f32, f32) {
10097    // SAFETY: callers uphold slice-length contracts (see dot_q4_row_avx2).
10098    unsafe {
10099        use core::arch::x86_64::*;
10100        let lomask = _mm_set1_epi8(0x0F);
10101        let eight = _mm256_set1_epi8(8);
10102        let ones = _mm256_set1_epi16(1);
10103        let (mut acc1, mut acc2) = (0f32, 0f32);
10104        #[inline(always)]
10105        unsafe fn hsum(d: core::arch::x86_64::__m256i) -> i32 {
10106            unsafe {
10107                use core::arch::x86_64::*;
10108                let hi128 = _mm256_extracti128_si256::<1>(d);
10109                let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
10110                let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
10111                let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
10112                _mm_cvtsi128_si32(s32)
10113            }
10114        }
10115        for gi in 0..gpr {
10116            let g = g0 + gi;
10117            let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
10118            let b = _mm_loadu_si128(packed.as_ptr().add(g * 16) as *const __m128i);
10119            let lo = _mm_and_si128(b, lomask);
10120            let hi = _mm_and_si128(_mm_srli_epi16::<4>(b), lomask);
10121            let w = _mm256_sub_epi8(
10122                _mm256_set_m128i(_mm_unpackhi_epi8(lo, hi), _mm_unpacklo_epi8(lo, hi)),
10123                eight,
10124            );
10125            let aw = _mm256_abs_epi8(w);
10126            let x1 = _mm256_loadu_si256(xq1.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
10127            let x2 = _mm256_loadu_si256(xq2.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
10128            let d1 = _mm256_madd_epi16(_mm256_maddubs_epi16(aw, _mm256_sign_epi8(x1, w)), ones);
10129            let d2 = _mm256_madd_epi16(_mm256_maddubs_epi16(aw, _mm256_sign_epi8(x2, w)), ones);
10130            acc1 += hsum(d1) as f32 * s;
10131            acc2 += hsum(d2) as f32 * s;
10132        }
10133        (acc1, acc2)
10134    }
10135}
10136
10137/// One q8 row range via AVX2 (x86 mirror of `q8_range_sdot`).
10138#[cfg(target_arch = "x86_64")]
10139fn q8_range_avx2(
10140    q: &[u8],
10141    row_scale: &[f32],
10142    act: &SplitAct,
10143    cols: usize,
10144    out_addr: SendMut,
10145    start: usize,
10146    end: usize,
10147) {
10148    for o in start..end {
10149        let v = row_dot_avx2(&q[o * cols..(o + 1) * cols], act) * row_scale[o];
10150        // SAFETY: disjoint row ranges per worker.
10151        unsafe { *out_addr.at(o) = v };
10152    }
10153}
10154
10155/// Two-input q8 row range via AVX2 (x86 mirror of `q8_range2_sdot`).
10156#[cfg(target_arch = "x86_64")]
10157#[allow(clippy::too_many_arguments)]
10158fn q8_range2_avx2(
10159    q: &[u8],
10160    row_scale: &[f32],
10161    a1: &SplitAct,
10162    a2: &SplitAct,
10163    cols: usize,
10164    p1: SendMut,
10165    p2: SendMut,
10166    start: usize,
10167    end: usize,
10168) {
10169    for o in start..end {
10170        let row = &q[o * cols..(o + 1) * cols];
10171        // SAFETY: disjoint row ranges per worker.
10172        unsafe {
10173            *p1.at(o) = row_dot_avx2(row, a1) * row_scale[o];
10174            *p2.at(o) = row_dot_avx2(row, a2) * row_scale[o];
10175        }
10176    }
10177}
10178
10179// ───────────────────── A8W8 SDOT path (port of vmfcore, ×1.78 decode) ─────────────────────
10180
10181/// ARMv8.6 i8mm (smmla): 32 int8 MACs per instruction vs sdot's 16 —
10182/// yet MEASURED 2.4× SLOWER than the blocked sdot on Apple silicon
10183/// (108 vs 264 GF/s): the on-the-fly vcombine packing and the two-
10184/// accumulator dependency chain swamp the MAC advantage, and Apple's
10185/// four SIMD pipes already keep sdot fed. OPT-IN (CMF_I8MM=1) for
10186/// field trials on Cortex-A710/X-class parts with two pipes, where the
10187/// balance may differ; a pre-interleaved weight layout (repack infra)
10188/// is the known path if it ever earns its keep.
10189#[cfg(target_arch = "aarch64")]
10190fn i8mm_enabled() -> bool {
10191    static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
10192    *ON.get_or_init(|| {
10193        std::env::var("CMF_I8MM").map(|v| v == "1").unwrap_or(false)
10194            && std::arch::is_aarch64_feature_detected!("i8mm")
10195    })
10196}
10197
10198/// SDOT enabled? Default ON when the CPU has ARMv8.2 dotprod;
10199/// `CMF_SDOT=0` disables (falls back to i8×f32 NEON).
10200/// (On non-ARM release builds only the test tolerance switch calls it.)
10201#[cfg_attr(not(target_arch = "aarch64"), allow(dead_code))]
10202fn sdot_enabled() -> bool {
10203    if FLOAT_ACTIVATIONS.get() {
10204        return false;
10205    }
10206    use std::sync::OnceLock;
10207    static ON: OnceLock<bool> = OnceLock::new();
10208    *ON.get_or_init(|| {
10209        let want = std::env::var("CMF_SDOT").map(|v| v != "0").unwrap_or(true);
10210        if !want {
10211            return false;
10212        }
10213
10214        #[cfg(target_arch = "aarch64")]
10215        {
10216            if std::arch::is_aarch64_feature_detected!("dotprod") {
10217                return true;
10218            }
10219            #[cfg(target_os = "android")]
10220            {
10221                if let Ok(cpuinfo) = std::fs::read_to_string("/proc/cpuinfo") {
10222                    if cpuinfo.lines().any(|l| {
10223                        (l.starts_with("Features") || l.starts_with("features"))
10224                            && l.contains("asimddp")
10225                    }) {
10226                        return true;
10227                    }
10228                }
10229            }
10230            false
10231        }
10232        #[cfg(not(target_arch = "aarch64"))]
10233        {
10234            false
10235        }
10236    })
10237}
10238
10239/// Two-field activation split (≡ vmfcore `q8_split_prep`): outlier
10240/// channels (>8·rms) are computed exactly in f32; the bulk (outliers
10241/// zeroed → clean absmax) goes through int8 SDOT. Computed ONCE per
10242/// matvec, shared by all rows/workers.
10243struct SplitAct {
10244    xq: Vec<i8>,
10245    sx: f32,
10246    outliers: Vec<(usize, f32)>,
10247    /// Σ xq — the VNNI bias-trick correction (`(w+128)·x` sums need
10248    /// `−128·Σx`); one i32 per split, computed once per matvec.
10249    #[cfg_attr(not(target_arch = "x86_64"), allow(dead_code))]
10250    xsum: i32,
10251}
10252
10253thread_local! {
10254    /// Recycled xq buffers: split_act runs for every matvec (~200/token)
10255    /// and its hidden-size allocation was steady-state heap churn.
10256    static XQ_FREE: std::cell::RefCell<Vec<Vec<i8>>> =
10257        const { std::cell::RefCell::new(Vec::new()) };
10258}
10259
10260impl Drop for SplitAct {
10261    fn drop(&mut self) {
10262        let buf = std::mem::take(&mut self.xq);
10263        if buf.capacity() > 0 {
10264            XQ_FREE.with(|f| {
10265                let mut f = f.borrow_mut();
10266                if f.len() < 16 {
10267                    f.push(buf);
10268                }
10269            });
10270        }
10271    }
10272}
10273
10274thread_local! {
10275    /// One scratch row per WORKER, kept for the life of the thread.
10276    ///
10277    /// The kernels take a row of group scales per dispatch, and a fresh
10278    /// `vec![0f32; gpr]` inside the closure is one allocation per worker per
10279    /// dispatch — on the release checkpoint about six thousand a token, a
10280    /// quarter of everything the benchmark counts.
10281    static KROW: std::cell::RefCell<Vec<f32>> = const { std::cell::RefCell::new(Vec::new()) };
10282}
10283
10284/// Borrow `n` floats of the calling worker's scratch. Nothing inside a
10285/// kernel body borrows it again, which is what keeps the RefCell honest.
10286#[inline]
10287fn with_krow<R>(n: usize, f: impl FnOnce(&mut [f32]) -> R) -> R {
10288    KROW.with(|s| {
10289        let mut b = s.borrow_mut();
10290        if b.len() < n {
10291            b.resize(n, 0.0);
10292        }
10293        f(&mut b[..n])
10294    })
10295}
10296
10297/// `t.round().clamp(-127.0, 127.0) as i8`, bit for bit, without the libm
10298/// call. On baseline x86-64 (no SSE4.1 `roundps`) `f32::round` is a
10299/// function call per element, and split_act runs it over every hidden
10300/// state before every matvec: measured 27 us a call on a 2048-wide
10301/// activation on an EPYC 7763 — 5.4 ms of a 55 ms decode token, all of
10302/// it on the caller's thread while thirty workers wait. Clamping first is
10303/// equivalent (round is monotonic and ±127 are integers), and after the
10304/// clamp `t - trunc(t)` is exact, so the half-away-from-zero decision is
10305/// the one `round` makes. NaN clamps to NaN and converts to 0, as before.
10306/// The loop vectorizes (cvttps2dq + compare/select).
10307#[inline(always)]
10308fn q8_round(t: f32) -> i8 {
10309    let t = t.clamp(-127.0, 127.0);
10310    let i = t as i32;
10311    let f = t - i as f32;
10312    let r = if f >= 0.5 {
10313        i + 1
10314    } else if f <= -0.5 {
10315        i - 1
10316    } else {
10317        i
10318    };
10319    r as i8
10320}
10321
10322fn split_act(x: &[f32]) -> SplitAct {
10323    let _prof = crate::cpuprof::time(crate::cpuprof::Slot::SplitAct);
10324    let n = x.len();
10325    let rms = (x.iter().map(|&v| (v * v) as f64).sum::<f64>() / n.max(1) as f64).sqrt() as f32;
10326    let thr = 8.0 * rms;
10327    // One pass: collect outliers and the bulk absmax (outliers excluded —
10328    // identical to the old zero-then-fold over a copied buffer, minus the
10329    // full-vector copy).
10330    let mut outliers: Vec<(usize, f32)> = Vec::new();
10331    let mut amax = 0f32;
10332    for (j, &v) in x.iter().enumerate() {
10333        let a = v.abs();
10334        if a > thr {
10335            outliers.push((j, v));
10336        } else if a > amax {
10337            amax = a;
10338        }
10339    }
10340    let sx = if amax > 0.0 { amax / 127.0 } else { 1.0 };
10341    let inv = 1.0 / sx;
10342    let mut xq = XQ_FREE.with(|f| f.borrow_mut().pop()).unwrap_or_default();
10343    xq.clear();
10344    xq.reserve(n);
10345    if outliers.is_empty() {
10346        xq.extend(
10347            x.iter()
10348                .map(|&v| q8_round(v * inv)),
10349        );
10350    } else {
10351        // Outlier slots quantize to 0 (their exact term is added later).
10352        xq.extend(x.iter().map(|&v| {
10353            if v.abs() > thr {
10354                0
10355            } else {
10356                q8_round(v * inv)
10357            }
10358        }));
10359    }
10360    let xsum = xq.iter().map(|&v| v as i32).sum();
10361    SplitAct {
10362        xq,
10363        sx,
10364        outliers,
10365        xsum,
10366    }
10367}
10368
10369fn split_act_q8_2f(x: &[f32], col: &[f32]) -> SplitAct {
10370    let _prof = crate::cpuprof::time(crate::cpuprof::Slot::SplitAct);
10371    let n = x.len();
10372    let rms = (x
10373        .iter()
10374        .zip(col)
10375        .map(|(&a, &c)| {
10376            let v = a * c;
10377            (v * v) as f64
10378        })
10379        .sum::<f64>()
10380        / n.max(1) as f64)
10381        .sqrt() as f32;
10382    let thr = 8.0 * rms;
10383
10384    let mut outliers = Vec::new();
10385    let mut amax = 0f32;
10386    for (j, (&a, &c)) in x.iter().zip(col).enumerate() {
10387        let v = a * c;
10388        let s = v.abs();
10389        if s > thr {
10390            outliers.push((j, v));
10391        } else if s > amax {
10392            amax = s;
10393        }
10394    }
10395
10396    let sx = if amax > 0.0 { amax / 127.0 } else { 1.0 };
10397    let inv = 1.0 / sx;
10398    let mut xq = XQ_FREE.with(|f| f.borrow_mut().pop()).unwrap_or_default();
10399    xq.clear();
10400    xq.reserve(n);
10401    if outliers.is_empty() {
10402        xq.extend(
10403            x.iter()
10404                .zip(col)
10405                .map(|(&a, &c)| q8_round((a * c) * inv)),
10406        );
10407    } else {
10408        xq.extend(x.iter().zip(col).map(|(&a, &c)| {
10409            let v = a * c;
10410            if v.abs() > thr {
10411                0
10412            } else {
10413                q8_round(v * inv)
10414            }
10415        }));
10416    }
10417    let xsum = xq.iter().map(|&v| v as i32).sum();
10418    SplitAct {
10419        xq,
10420        sx,
10421        outliers,
10422        xsum,
10423    }
10424}
10425
10426/// int8(weight)·int8(activation) → i32 via `sdot` (inline asm — the
10427/// vdotq intrinsic is unstable; port of vmfcore `dot_i8_sdot`).
10428#[cfg(target_arch = "aarch64")]
10429#[target_feature(enable = "neon,dotprod")]
10430unsafe fn dot_i8_sdot(w: &[u8], xq: &[i8]) -> i32 {
10431    // SAFETY: callers uphold slice-length contracts (see call sites).
10432    unsafe {
10433        use core::arch::aarch64::*;
10434        use core::arch::asm;
10435        let wp = w.as_ptr() as *const i8;
10436        let n = w.len();
10437        let (mut a0, mut a1, mut a2, mut a3) = (
10438            vdupq_n_s32(0),
10439            vdupq_n_s32(0),
10440            vdupq_n_s32(0),
10441            vdupq_n_s32(0),
10442        );
10443        let mut i = 0;
10444        while i + 64 <= n {
10445            let (w0, x0) = (vld1q_s8(wp.add(i)), vld1q_s8(xq.as_ptr().add(i)));
10446            let (w1, x1) = (vld1q_s8(wp.add(i + 16)), vld1q_s8(xq.as_ptr().add(i + 16)));
10447            let (w2, x2) = (vld1q_s8(wp.add(i + 32)), vld1q_s8(xq.as_ptr().add(i + 32)));
10448            let (w3, x3) = (vld1q_s8(wp.add(i + 48)), vld1q_s8(xq.as_ptr().add(i + 48)));
10449            asm!(
10450                "sdot {a0:v}.4s, {w0:v}.16b, {x0:v}.16b",
10451                "sdot {a1:v}.4s, {w1:v}.16b, {x1:v}.16b",
10452                "sdot {a2:v}.4s, {w2:v}.16b, {x2:v}.16b",
10453                "sdot {a3:v}.4s, {w3:v}.16b, {x3:v}.16b",
10454                a0 = inout(vreg) a0, a1 = inout(vreg) a1, a2 = inout(vreg) a2, a3 = inout(vreg) a3,
10455                w0 = in(vreg) w0, x0 = in(vreg) x0, w1 = in(vreg) w1, x1 = in(vreg) x1,
10456                w2 = in(vreg) w2, x2 = in(vreg) x2, w3 = in(vreg) w3, x3 = in(vreg) x3,
10457                options(pure, nomem, nostack),
10458            );
10459            i += 64;
10460        }
10461        while i + 16 <= n {
10462            let (wv, xv) = (vld1q_s8(wp.add(i)), vld1q_s8(xq.as_ptr().add(i)));
10463            asm!("sdot {a:v}.4s, {w:v}.16b, {x:v}.16b",
10464                 a = inout(vreg) a0, w = in(vreg) wv, x = in(vreg) xv, options(pure, nomem, nostack));
10465            i += 16;
10466        }
10467        let mut s = vaddvq_s32(vaddq_s32(vaddq_s32(a0, a1), vaddq_s32(a2, a3)));
10468        while i < n {
10469            s += (*wp.add(i)) as i32 * xq[i] as i32;
10470            i += 1;
10471        }
10472        s
10473    }
10474}
10475
10476/// Row-blocked SDOT: 4 output rows per pass — the activation chunk is
10477/// loaded once and reused, 4 independent accumulators hide sdot latency
10478/// (port of vmfcore `dot_i8_sdot_4rows`).
10479#[cfg(target_arch = "aarch64")]
10480#[target_feature(enable = "neon,dotprod")]
10481unsafe fn dot_i8_sdot_4rows(w0: &[u8], w1: &[u8], w2: &[u8], w3: &[u8], xq: &[i8]) -> [i32; 4] {
10482    // SAFETY: callers uphold slice-length contracts (see call sites).
10483    unsafe {
10484        use core::arch::aarch64::*;
10485        use core::arch::asm;
10486        let n = xq.len();
10487        let px = xq.as_ptr();
10488        let (p0, p1, p2, p3) = (
10489            w0.as_ptr() as *const i8,
10490            w1.as_ptr() as *const i8,
10491            w2.as_ptr() as *const i8,
10492            w3.as_ptr() as *const i8,
10493        );
10494        let (mut a0, mut a1, mut a2, mut a3) = (
10495            vdupq_n_s32(0),
10496            vdupq_n_s32(0),
10497            vdupq_n_s32(0),
10498            vdupq_n_s32(0),
10499        );
10500        let mut i = 0;
10501        while i + 16 <= n {
10502            let x = vld1q_s8(px.add(i));
10503            let v0 = vld1q_s8(p0.add(i));
10504            let v1 = vld1q_s8(p1.add(i));
10505            let v2 = vld1q_s8(p2.add(i));
10506            let v3 = vld1q_s8(p3.add(i));
10507            asm!(
10508                "sdot {a0:v}.4s, {v0:v}.16b, {x:v}.16b",
10509                "sdot {a1:v}.4s, {v1:v}.16b, {x:v}.16b",
10510                "sdot {a2:v}.4s, {v2:v}.16b, {x:v}.16b",
10511                "sdot {a3:v}.4s, {v3:v}.16b, {x:v}.16b",
10512                a0 = inout(vreg) a0, a1 = inout(vreg) a1, a2 = inout(vreg) a2, a3 = inout(vreg) a3,
10513                v0 = in(vreg) v0, v1 = in(vreg) v1, v2 = in(vreg) v2, v3 = in(vreg) v3, x = in(vreg) x,
10514                options(pure, nomem, nostack),
10515            );
10516            i += 16;
10517        }
10518        let mut r = [
10519            vaddvq_s32(a0),
10520            vaddvq_s32(a1),
10521            vaddvq_s32(a2),
10522            vaddvq_s32(a3),
10523        ];
10524        while i < n {
10525            let xi = *px.add(i) as i32;
10526            r[0] += (*p0.add(i)) as i32 * xi;
10527            r[1] += (*p1.add(i)) as i32 * xi;
10528            r[2] += (*p2.add(i)) as i32 * xi;
10529            r[3] += (*p3.add(i)) as i32 * xi;
10530            i += 1;
10531        }
10532        r
10533    }
10534}
10535
10536/// 4 interleaved rows in one pass: the repacked group is [r0[c], r1[c],
10537/// r2[c], r3[c]] per 16-byte chunk, so each iteration reads ONE 64-byte
10538/// line plus the shared activation chunk — a single sequential weight
10539/// stream per worker. Per-row accumulation is the same one-accumulator
10540/// scheme as `dot_i8_sdot_4rows`; integer sums are exact, so outputs
10541/// are bit-identical to the mmap-layout kernel.
10542#[cfg(target_arch = "aarch64")]
10543#[target_feature(enable = "neon,dotprod")]
10544unsafe fn dot_i8_sdot_4rows_il(g: &[u8], xq: &[i8]) -> [i32; 4] {
10545    // SAFETY: callers uphold slice-length contracts (g.len() == 4·n,
10546    // n % 16 == 0 — guaranteed by the repack gate).
10547    unsafe {
10548        use core::arch::aarch64::*;
10549        use core::arch::asm;
10550        let n = xq.len();
10551        let px = xq.as_ptr();
10552        let pg = g.as_ptr() as *const i8;
10553        let (mut a0, mut a1, mut a2, mut a3) = (
10554            vdupq_n_s32(0),
10555            vdupq_n_s32(0),
10556            vdupq_n_s32(0),
10557            vdupq_n_s32(0),
10558        );
10559        let mut i = 0;
10560        while i + 16 <= n {
10561            let x = vld1q_s8(px.add(i));
10562            let base = pg.add(4 * i);
10563            let v0 = vld1q_s8(base);
10564            let v1 = vld1q_s8(base.add(16));
10565            let v2 = vld1q_s8(base.add(32));
10566            let v3 = vld1q_s8(base.add(48));
10567            asm!(
10568                "sdot {a0:v}.4s, {v0:v}.16b, {x:v}.16b",
10569                "sdot {a1:v}.4s, {v1:v}.16b, {x:v}.16b",
10570                "sdot {a2:v}.4s, {v2:v}.16b, {x:v}.16b",
10571                "sdot {a3:v}.4s, {v3:v}.16b, {x:v}.16b",
10572                a0 = inout(vreg) a0, a1 = inout(vreg) a1, a2 = inout(vreg) a2, a3 = inout(vreg) a3,
10573                v0 = in(vreg) v0, v1 = in(vreg) v1, v2 = in(vreg) v2, v3 = in(vreg) v3, x = in(vreg) x,
10574                options(pure, nomem, nostack),
10575            );
10576            i += 16;
10577        }
10578        [
10579            vaddvq_s32(a0),
10580            vaddvq_s32(a1),
10581            vaddvq_s32(a2),
10582            vaddvq_s32(a3),
10583        ]
10584    }
10585}
10586
10587/// One q8 row range via SDOT (4-row blocks + tail) — the body of
10588/// `qmatvec`'s hot loop, extracted so multi-matrix jobs can drive the
10589/// SAME kernel for several tensors under one pool dispatch. `rep` — the
10590/// load-time interleaved repack (empty = mmap layout only); rows outside
10591/// full 4-row groups always come from the mmap layout.
10592#[cfg(target_arch = "aarch64")]
10593fn q8_range_sdot(
10594    q: &[u8],
10595    rep: &[u8],
10596    row_scale: &[f32],
10597    act: &SplitAct,
10598    cols: usize,
10599    out_addr: SendMut,
10600    start: usize,
10601    end: usize,
10602) {
10603    let mut o = start;
10604    // Leading rows to the group boundary (repack path only): the pool
10605    // splits row ranges arbitrarily, groups are absolute.
10606    if !rep.is_empty() {
10607        while o < end && o % 4 != 0 {
10608            let v = row_dot_sdot(&q[o * cols..(o + 1) * cols], act) * row_scale[o];
10609            unsafe { *out_addr.at(o) = v };
10610            o += 1;
10611        }
10612    }
10613    while o + 4 <= end {
10614        let r = if rep.is_empty() {
10615            unsafe {
10616                dot_i8_sdot_4rows(
10617                    &q[o * cols..(o + 1) * cols],
10618                    &q[(o + 1) * cols..(o + 2) * cols],
10619                    &q[(o + 2) * cols..(o + 3) * cols],
10620                    &q[(o + 3) * cols..(o + 4) * cols],
10621                    &act.xq,
10622                )
10623            }
10624        } else {
10625            unsafe { dot_i8_sdot_4rows_il(&rep[o * cols..(o + 4) * cols], &act.xq) }
10626        };
10627        for k in 0..4 {
10628            let mut acc = r[k] as f32 * act.sx;
10629            for &(j, xv) in &act.outliers {
10630                acc += (q[(o + k) * cols + j] as i8) as f32 * xv;
10631            }
10632            // SAFETY: disjoint row ranges per worker.
10633            unsafe { *out_addr.at(o + k) = acc * row_scale[o + k] };
10634        }
10635        o += 4;
10636    }
10637    while o < end {
10638        let v = row_dot_sdot(&q[o * cols..(o + 1) * cols], act) * row_scale[o];
10639        unsafe { *out_addr.at(o) = v };
10640        o += 1;
10641    }
10642}
10643
10644/// Two-input q8 row range via SDOT — `qmatvec2`'s hot loop, extracted
10645/// for the fused pair multi-matrix job (`matvec2_many`).
10646#[cfg(target_arch = "aarch64")]
10647#[allow(clippy::too_many_arguments)]
10648fn q8_range2_sdot(
10649    q: &[u8],
10650    row_scale: &[f32],
10651    a1: &SplitAct,
10652    a2: &SplitAct,
10653    cols: usize,
10654    p1: SendMut,
10655    p2: SendMut,
10656    start: usize,
10657    end: usize,
10658) {
10659    for o in start..end {
10660        let row = &q[o * cols..(o + 1) * cols];
10661        // SAFETY: disjoint row ranges per worker.
10662        unsafe {
10663            *p1.at(o) = row_dot_sdot(row, a1) * row_scale[o];
10664            *p2.at(o) = row_dot_sdot(row, a2) * row_scale[o];
10665        }
10666    }
10667}
10668
10669/// Two-input q8 row range, f32 kernel (non-SDOT) — same extraction.
10670#[allow(clippy::too_many_arguments)]
10671fn q8_range2_f32(
10672    q: &[u8],
10673    row_scale: &[f32],
10674    x1: &[f32],
10675    x2: &[f32],
10676    cols: usize,
10677    p1: SendMut,
10678    p2: SendMut,
10679    start: usize,
10680    end: usize,
10681) {
10682    for o in start..end {
10683        let row = &q[o * cols..(o + 1) * cols];
10684        // SAFETY: disjoint row ranges per worker.
10685        unsafe {
10686            *p1.at(o) = dot_i8_f32(row, x1) * row_scale[o];
10687            *p2.at(o) = dot_i8_f32(row, x2) * row_scale[o];
10688        }
10689    }
10690}
10691
10692/// Scalar/NEON-f32 q8 row range (non-SDOT platforms) — same extraction.
10693fn q8_range_f32(
10694    q: &[u8],
10695    row_scale: &[f32],
10696    xs: &[f32],
10697    cols: usize,
10698    out_addr: SendMut,
10699    start: usize,
10700    end: usize,
10701) {
10702    for o in start..end {
10703        let v = dot_i8_f32(&q[o * cols..(o + 1) * cols], xs) * row_scale[o];
10704        // SAFETY: disjoint row ranges per worker.
10705        unsafe { *out_addr.at(o) = v };
10706    }
10707}
10708
10709/// One q8 row against a split activation, portable: the per-arch fast
10710/// dots where they exist, the exact scalar loop elsewhere. The scalar
10711/// arm is also the test oracle for both fast arms.
10712#[inline]
10713fn q8_row_dot(row: &[u8], act: &SplitAct) -> f32 {
10714    #[cfg(target_arch = "aarch64")]
10715    return row_dot_sdot(row, act);
10716    #[cfg(target_arch = "x86_64")]
10717    return row_dot_avx2(row, act);
10718    #[allow(unreachable_code)]
10719    q8_row_dot_scalar(row, act)
10720}
10721
10722#[allow(dead_code)]
10723fn q8_row_dot_scalar(row: &[u8], act: &SplitAct) -> f32 {
10724    let mut acc = 0i32;
10725    for (k, &b) in row.iter().enumerate() {
10726        acc += (b as i8) as i32 * act.xq[k] as i32;
10727    }
10728    let mut acc = acc as f32 * act.sx;
10729    for &(j, xv) in &act.outliers {
10730        acc += (row[j] as i8) as f32 * xv;
10731    }
10732    acc
10733}
10734
10735/// SDOT row dot with exact outlier correction:
10736/// `dot = sdot(w, xq)·sx + Σ_outl w[j]·x[j]` (then × row_scale by caller).
10737#[cfg(target_arch = "aarch64")]
10738#[inline]
10739fn row_dot_sdot(row: &[u8], act: &SplitAct) -> f32 {
10740    let mut acc = unsafe { dot_i8_sdot(row, &act.xq) } as f32 * act.sx;
10741    for &(j, xv) in &act.outliers {
10742        acc += (row[j] as i8) as f32 * xv;
10743    }
10744    acc
10745}
10746
10747/// One q4 row via SDOT: each 32-group's nibbles unpack to centered i8
10748/// (nib−8 ∈ [−8,7]), int8×int8 `sdot` against the pre-quantized
10749/// activation group, × the group's f16 scale. Returns Σ_g dot_g·s_g;
10750/// the caller multiplies by the activation scale and adds the exact
10751/// outlier terms (port of vmfcore `dot_q4_block_sdot`, +23% measured).
10752/// Nibble order matches the writer: element 2k = low nibble, 2k+1 = high
10753/// → zip(lo,hi) restores flat order.
10754#[cfg(target_arch = "aarch64")]
10755#[target_feature(enable = "neon,dotprod")]
10756unsafe fn dot_q4_row_sdot(packed: &[u8], scales: &[u8], g0: usize, gpr: usize, xq: &[i8]) -> f32 {
10757    // SAFETY: callers uphold slice-length contracts (16 packed bytes and
10758    // 2 scale bytes per group; xq.len() == gpr·GROUP_SIZE).
10759    unsafe {
10760        use core::arch::aarch64::*;
10761        use core::arch::asm;
10762        let lomask = vdupq_n_u8(0x0F);
10763        let eight = vdupq_n_s8(8);
10764        let mut acc = 0f32;
10765        for gi in 0..gpr {
10766            let g = g0 + gi;
10767            let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
10768            let b = vld1q_u8(packed.as_ptr().add(g * 16));
10769            let lo = vandq_u8(b, lomask);
10770            let hi = vshrq_n_u8::<4>(b);
10771            let e0 = vsubq_s8(vreinterpretq_s8_u8(vzip1q_u8(lo, hi)), eight);
10772            let e1 = vsubq_s8(vreinterpretq_s8_u8(vzip2q_u8(lo, hi)), eight);
10773            let x0 = vld1q_s8(xq.as_ptr().add(gi * GROUP_SIZE));
10774            let x1 = vld1q_s8(xq.as_ptr().add(gi * GROUP_SIZE + 16));
10775            let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
10776            asm!(
10777                "sdot {a0:v}.4s, {e0:v}.16b, {x0:v}.16b",
10778                "sdot {a1:v}.4s, {e1:v}.16b, {x1:v}.16b",
10779                a0 = inout(vreg) a0, a1 = inout(vreg) a1,
10780                e0 = in(vreg) e0, x0 = in(vreg) x0, e1 = in(vreg) e1, x1 = in(vreg) x1,
10781                options(pure, nomem, nostack),
10782            );
10783            acc += vaddvq_s32(vaddq_s32(a0, a1)) as f32 * s;
10784        }
10785        acc
10786    }
10787}
10788
10789/// Two-activation q4 row via SDOT: the nibble unpack (the expensive
10790/// part) happens ONCE per group; both pre-quantized activations are
10791/// dotted against the same centered i8 registers. Per-lane math matches
10792/// `dot_q4_row_sdot` exactly.
10793#[cfg(target_arch = "aarch64")]
10794#[target_feature(enable = "neon,dotprod")]
10795unsafe fn dot_q4_row_sdot2(
10796    packed: &[u8],
10797    scales: &[u8],
10798    g0: usize,
10799    gpr: usize,
10800    xq1: &[i8],
10801    xq2: &[i8],
10802) -> (f32, f32) {
10803    // SAFETY: callers uphold slice-length contracts (16 packed bytes and
10804    // 2 scale bytes per group; xq*.len() == gpr·GROUP_SIZE).
10805    unsafe {
10806        use core::arch::aarch64::*;
10807        use core::arch::asm;
10808        let lomask = vdupq_n_u8(0x0F);
10809        let eight = vdupq_n_s8(8);
10810        let (mut acc1, mut acc2) = (0f32, 0f32);
10811        for gi in 0..gpr {
10812            let g = g0 + gi;
10813            let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
10814            let b = vld1q_u8(packed.as_ptr().add(g * 16));
10815            let lo = vandq_u8(b, lomask);
10816            let hi = vshrq_n_u8::<4>(b);
10817            let e0 = vsubq_s8(vreinterpretq_s8_u8(vzip1q_u8(lo, hi)), eight);
10818            let e1 = vsubq_s8(vreinterpretq_s8_u8(vzip2q_u8(lo, hi)), eight);
10819            let x10 = vld1q_s8(xq1.as_ptr().add(gi * GROUP_SIZE));
10820            let x11 = vld1q_s8(xq1.as_ptr().add(gi * GROUP_SIZE + 16));
10821            let x20 = vld1q_s8(xq2.as_ptr().add(gi * GROUP_SIZE));
10822            let x21 = vld1q_s8(xq2.as_ptr().add(gi * GROUP_SIZE + 16));
10823            let (mut a0, mut a1, mut b0, mut b1) = (
10824                vdupq_n_s32(0),
10825                vdupq_n_s32(0),
10826                vdupq_n_s32(0),
10827                vdupq_n_s32(0),
10828            );
10829            asm!(
10830                "sdot {a0:v}.4s, {e0:v}.16b, {x10:v}.16b",
10831                "sdot {a1:v}.4s, {e1:v}.16b, {x11:v}.16b",
10832                "sdot {b0:v}.4s, {e0:v}.16b, {x20:v}.16b",
10833                "sdot {b1:v}.4s, {e1:v}.16b, {x21:v}.16b",
10834                a0 = inout(vreg) a0, a1 = inout(vreg) a1,
10835                b0 = inout(vreg) b0, b1 = inout(vreg) b1,
10836                e0 = in(vreg) e0, e1 = in(vreg) e1,
10837                x10 = in(vreg) x10, x11 = in(vreg) x11,
10838                x20 = in(vreg) x20, x21 = in(vreg) x21,
10839                options(pure, nomem, nostack),
10840            );
10841            acc1 += vaddvq_s32(vaddq_s32(a0, a1)) as f32 * s;
10842            acc2 += vaddvq_s32(vaddq_s32(b0, b1)) as f32 * s;
10843        }
10844        (acc1, acc2)
10845    }
10846}
10847
10848// ───────────────────── fused int8 kernels ─────────────────────
10849
10850/// `acc += w · row` where the row is centered i8 — NEON widen+fma on
10851/// aarch64, scalar elsewhere. The KV-cache q8 value path rides on this.
10852#[inline]
10853pub(crate) fn axpy_i8_f32(acc: &mut [f32], row: &[i8], w: f32) {
10854    #[cfg(target_arch = "aarch64")]
10855    unsafe {
10856        return axpy_i8_f32_neon(acc, row, w);
10857    }
10858    #[cfg(target_arch = "x86_64")]
10859    if avx2_enabled() {
10860        return unsafe { axpy_i8_f32_avx2(acc, row, w) };
10861    }
10862    #[allow(unreachable_code)]
10863    {
10864        for (a, &b) in acc.iter_mut().zip(row) {
10865            *a += w * b as f32;
10866        }
10867    }
10868}
10869
10870/// i8→f32 axpy via AVX2/FMA (x86 mirror of `axpy_i8_f32_neon`).
10871#[cfg(target_arch = "x86_64")]
10872#[target_feature(enable = "avx2,fma")]
10873unsafe fn axpy_i8_f32_avx2(acc: &mut [f32], row: &[i8], w: f32) {
10874    // SAFETY: callers uphold slice-length contracts (see call sites).
10875    unsafe {
10876        use core::arch::x86_64::*;
10877        let n = acc.len().min(row.len());
10878        let ap = acc.as_mut_ptr();
10879        let rp = row.as_ptr();
10880        let wv = _mm256_set1_ps(w);
10881        let mut j = 0usize;
10882        while j + 16 <= n {
10883            let rb = _mm_loadu_si128(rp.add(j) as *const __m128i);
10884            let lo = _mm256_cvtepi32_ps(_mm256_cvtepi8_epi32(rb));
10885            let hi = _mm256_cvtepi32_ps(_mm256_cvtepi8_epi32(_mm_srli_si128::<8>(rb)));
10886            let v0 = _mm256_fmadd_ps(wv, lo, _mm256_loadu_ps(ap.add(j)));
10887            let v1 = _mm256_fmadd_ps(wv, hi, _mm256_loadu_ps(ap.add(j + 8)));
10888            _mm256_storeu_ps(ap.add(j), v0);
10889            _mm256_storeu_ps(ap.add(j + 8), v1);
10890            j += 16;
10891        }
10892        while j < n {
10893            *ap.add(j) += w * (*rp.add(j)) as f32;
10894            j += 1;
10895        }
10896    }
10897}
10898
10899#[cfg(target_arch = "aarch64")]
10900#[target_feature(enable = "neon")]
10901unsafe fn axpy_i8_f32_neon(acc: &mut [f32], row: &[i8], w: f32) {
10902    // SAFETY: callers uphold slice-length contracts (see call sites).
10903    unsafe {
10904        use core::arch::aarch64::*;
10905        let n = acc.len().min(row.len());
10906        let ap = acc.as_mut_ptr();
10907        let rp = row.as_ptr();
10908        let wv = vdupq_n_f32(w);
10909        let mut j = 0usize;
10910        while j + 16 <= n {
10911            let rb = vld1q_s8(rp.add(j));
10912            let lo = vmovl_s8(vget_low_s8(rb));
10913            let hi = vmovl_s8(vget_high_s8(rb));
10914            for (off, half) in [(0, lo), (8, hi)] {
10915                let f0 = vcvtq_f32_s32(vmovl_s16(vget_low_s16(half)));
10916                let f1 = vcvtq_f32_s32(vmovl_s16(vget_high_s16(half)));
10917                let o = j + off;
10918                vst1q_f32(ap.add(o), vfmaq_f32(vld1q_f32(ap.add(o)), wv, f0));
10919                vst1q_f32(ap.add(o + 4), vfmaq_f32(vld1q_f32(ap.add(o + 4)), wv, f1));
10920            }
10921            j += 16;
10922        }
10923        while j < n {
10924            *ap.add(j) += w * (*rp.add(j)) as f32;
10925            j += 1;
10926        }
10927    }
10928}
10929
10930/// i8 row · f32 x. NEON on aarch64 (ported from vmfcore `dot_i8_f32_neon`,
10931/// ≈9× scalar), scalar elsewhere.
10932#[inline]
10933pub(crate) fn dot_i8_f32(w: &[u8], x: &[f32]) -> f32 {
10934    #[cfg(target_arch = "aarch64")]
10935    unsafe {
10936        return dot_i8_f32_neon(w, x);
10937    }
10938    #[cfg(target_arch = "x86_64")]
10939    if avx2_enabled() {
10940        return unsafe { dot_i8_f32_avx2(w, x) };
10941    }
10942    #[allow(unreachable_code)]
10943    {
10944        let mut sum = 0.0f32;
10945        for (j, &b) in w.iter().enumerate() {
10946            sum += (b as i8) as f32 * x[j];
10947        }
10948        sum
10949    }
10950}
10951
10952/// i8 row · (x ⊙ col_field) — the q8_2f row dot with the θ col-field
10953/// folded into the product (no prescaled copy of x). NEON on aarch64,
10954/// scalar elsewhere. Used by the active-neuron path `row_dot`.
10955#[inline]
10956fn dot_i8_col_f32(w: &[u8], x: &[f32], col: &[f32]) -> f32 {
10957    #[cfg(target_arch = "aarch64")]
10958    unsafe {
10959        return dot_i8_col_f32_neon(w, x, col);
10960    }
10961    #[allow(unreachable_code)]
10962    {
10963        let mut sum = 0.0f32;
10964        for (j, &b) in w.iter().enumerate() {
10965            sum += (b as i8) as f32 * x[j] * col[j];
10966        }
10967        sum
10968    }
10969}
10970
10971#[cfg(target_arch = "aarch64")]
10972#[target_feature(enable = "neon")]
10973unsafe fn dot_i8_col_f32_neon(w: &[u8], x: &[f32], col: &[f32]) -> f32 {
10974    // SAFETY: callers uphold slice-length contracts (see call sites).
10975    unsafe {
10976        use core::arch::aarch64::*;
10977        let n = x.len();
10978        let wp = w.as_ptr() as *const i8;
10979        let xp = x.as_ptr();
10980        let cp = col.as_ptr();
10981        let (mut a0, mut a1, mut a2, mut a3) = (
10982            vdupq_n_f32(0.0),
10983            vdupq_n_f32(0.0),
10984            vdupq_n_f32(0.0),
10985            vdupq_n_f32(0.0),
10986        );
10987        let mut j = 0usize;
10988        while j + 16 <= n {
10989            let wb = vld1q_s8(wp.add(j));
10990            let lo = vmovl_s8(vget_low_s8(wb));
10991            let hi = vmovl_s8(vget_high_s8(wb));
10992            let w0 = vcvtq_f32_s32(vmovl_s16(vget_low_s16(lo)));
10993            let w1 = vcvtq_f32_s32(vmovl_s16(vget_high_s16(lo)));
10994            let w2 = vcvtq_f32_s32(vmovl_s16(vget_low_s16(hi)));
10995            let w3 = vcvtq_f32_s32(vmovl_s16(vget_high_s16(hi)));
10996            a0 = vfmaq_f32(
10997                a0,
10998                w0,
10999                vmulq_f32(vld1q_f32(xp.add(j)), vld1q_f32(cp.add(j))),
11000            );
11001            a1 = vfmaq_f32(
11002                a1,
11003                w1,
11004                vmulq_f32(vld1q_f32(xp.add(j + 4)), vld1q_f32(cp.add(j + 4))),
11005            );
11006            a2 = vfmaq_f32(
11007                a2,
11008                w2,
11009                vmulq_f32(vld1q_f32(xp.add(j + 8)), vld1q_f32(cp.add(j + 8))),
11010            );
11011            a3 = vfmaq_f32(
11012                a3,
11013                w3,
11014                vmulq_f32(vld1q_f32(xp.add(j + 12)), vld1q_f32(cp.add(j + 12))),
11015            );
11016            j += 16;
11017        }
11018        let mut sum = vaddvq_f32(vaddq_f32(vaddq_f32(a0, a1), vaddq_f32(a2, a3)));
11019        while j < n {
11020            sum += (*wp.add(j)) as f32 * *xp.add(j) * *cp.add(j);
11021            j += 1;
11022        }
11023        sum
11024    }
11025}
11026
11027#[cfg(target_arch = "aarch64")]
11028#[target_feature(enable = "neon")]
11029unsafe fn dot_i8_f32_neon(w: &[u8], x: &[f32]) -> f32 {
11030    // SAFETY: callers uphold slice-length contracts (see call sites).
11031    unsafe {
11032        use core::arch::aarch64::*;
11033        let n = x.len();
11034        let wp = w.as_ptr() as *const i8;
11035        let xp = x.as_ptr();
11036        let (mut a0, mut a1, mut a2, mut a3) = (
11037            vdupq_n_f32(0.0),
11038            vdupq_n_f32(0.0),
11039            vdupq_n_f32(0.0),
11040            vdupq_n_f32(0.0),
11041        );
11042        let mut j = 0usize;
11043        while j + 16 <= n {
11044            let wb = vld1q_s8(wp.add(j));
11045            let lo = vmovl_s8(vget_low_s8(wb));
11046            let hi = vmovl_s8(vget_high_s8(wb));
11047            let w0 = vcvtq_f32_s32(vmovl_s16(vget_low_s16(lo)));
11048            let w1 = vcvtq_f32_s32(vmovl_s16(vget_high_s16(lo)));
11049            let w2 = vcvtq_f32_s32(vmovl_s16(vget_low_s16(hi)));
11050            let w3 = vcvtq_f32_s32(vmovl_s16(vget_high_s16(hi)));
11051            a0 = vfmaq_f32(a0, w0, vld1q_f32(xp.add(j)));
11052            a1 = vfmaq_f32(a1, w1, vld1q_f32(xp.add(j + 4)));
11053            a2 = vfmaq_f32(a2, w2, vld1q_f32(xp.add(j + 8)));
11054            a3 = vfmaq_f32(a3, w3, vld1q_f32(xp.add(j + 12)));
11055            j += 16;
11056        }
11057        let mut sum = vaddvq_f32(vaddq_f32(vaddq_f32(a0, a1), vaddq_f32(a2, a3)));
11058        while j < n {
11059            sum += (*wp.add(j)) as f32 * *xp.add(j);
11060            j += 1;
11061        }
11062        sum
11063    }
11064}
11065
11066#[allow(clippy::too_many_arguments)]
11067fn qmatvec(
11068    q: &[u8],
11069    rep: &[u8],
11070    row_scale: &[f32],
11071    x: &[f32],
11072    col_field: &[f32],
11073    dtype: TensorDtype,
11074    rows: usize,
11075    cols: usize,
11076    out: &mut [f32],
11077    pool: Option<&Pool>,
11078) {
11079    debug_assert_eq!(out.len(), rows);
11080    #[cfg(not(target_arch = "aarch64"))]
11081    let _ = rep;
11082
11083    #[cfg(target_arch = "aarch64")]
11084    if sdot_enabled() {
11085        let act = if dtype == TensorDtype::Q8_2f {
11086            split_act_q8_2f(x, col_field)
11087        } else {
11088            split_act(x)
11089        };
11090        let out_addr = SendMut(out.as_mut_ptr());
11091        let run_range = |start: usize, end: usize| {
11092            q8_range_sdot(q, rep, row_scale, &act, cols, out_addr, start, end)
11093        };
11094        match pool {
11095            Some(pool) if rows >= 256 => pool.run_rows(rows, &run_range),
11096            _ => run_range(0, rows),
11097        }
11098        return;
11099    }
11100    // x86 A8W8 via AVX2 maddubs — same quantized-activation contract as
11101    // the SDOT path (CMF_AVX2=0 keeps the exact i8×f32 loop).
11102    #[cfg(target_arch = "x86_64")]
11103    if avx2_a8w8_enabled() {
11104        let act = if dtype == TensorDtype::Q8_2f {
11105            split_act_q8_2f(x, col_field)
11106        } else {
11107            split_act(x)
11108        };
11109        let out_addr = SendMut(out.as_mut_ptr());
11110        let run_range = |start: usize, end: usize| {
11111            q8_range_avx2(q, row_scale, &act, cols, out_addr, start, end)
11112        };
11113        match pool {
11114            Some(pool) if rows >= 256 => pool.run_rows(rows, &run_range),
11115            _ => run_range(0, rows),
11116        }
11117        return;
11118    }
11119
11120    prescale_with(x, col_field, dtype, 1, |xs| {
11121        let out_addr = SendMut(out.as_mut_ptr());
11122        let run_range = move |start: usize, end: usize| {
11123            for o in start..end {
11124                let v = dot_i8_f32(&q[o * cols..(o + 1) * cols], xs) * row_scale[o];
11125                // SAFETY: disjoint row ranges per worker.
11126                unsafe { *out_addr.at(o) = v };
11127            }
11128        };
11129        match pool {
11130            Some(pool) if rows >= 256 => pool.run_rows(rows, &run_range),
11131            _ => run_range(0, rows),
11132        }
11133    });
11134}
11135
11136#[allow(clippy::too_many_arguments)]
11137fn qmatvec2(
11138    q: &[u8],
11139    row_scale: &[f32],
11140    x1: &[f32],
11141    x2: &[f32],
11142    col_field: &[f32],
11143    dtype: TensorDtype,
11144    rows: usize,
11145    cols: usize,
11146    o1: &mut [f32],
11147    o2: &mut [f32],
11148    pool: Option<&Pool>,
11149) {
11150    #[cfg(target_arch = "aarch64")]
11151    if sdot_enabled() {
11152        let a1s = if dtype == TensorDtype::Q8_2f {
11153            split_act_q8_2f(x1, col_field)
11154        } else {
11155            split_act(x1)
11156        };
11157        let a2s = if dtype == TensorDtype::Q8_2f {
11158            split_act_q8_2f(x2, col_field)
11159        } else {
11160            split_act(x2)
11161        };
11162        let p1 = SendMut(o1.as_mut_ptr());
11163        let p2 = SendMut(o2.as_mut_ptr());
11164        let run_range = |start: usize, end: usize| {
11165            q8_range2_sdot(q, row_scale, &a1s, &a2s, cols, p1, p2, start, end)
11166        };
11167        match pool {
11168            Some(pool) if rows >= 256 => pool.run_rows(rows, &run_range),
11169            _ => run_range(0, rows),
11170        }
11171        return;
11172    }
11173    #[cfg(target_arch = "x86_64")]
11174    if avx2_a8w8_enabled() {
11175        let a1s = if dtype == TensorDtype::Q8_2f {
11176            split_act_q8_2f(x1, col_field)
11177        } else {
11178            split_act(x1)
11179        };
11180        let a2s = if dtype == TensorDtype::Q8_2f {
11181            split_act_q8_2f(x2, col_field)
11182        } else {
11183            split_act(x2)
11184        };
11185        let p1 = SendMut(o1.as_mut_ptr());
11186        let p2 = SendMut(o2.as_mut_ptr());
11187        let run_range = |start: usize, end: usize| {
11188            q8_range2_avx2(q, row_scale, &a1s, &a2s, cols, p1, p2, start, end)
11189        };
11190        match pool {
11191            Some(pool) if rows >= 256 => pool.run_rows(rows, &run_range),
11192            _ => run_range(0, rows),
11193        }
11194        return;
11195    }
11196
11197    prescale_with(x1, col_field, dtype, 1, |x1s| {
11198        prescale_with(x2, col_field, dtype, 2, |x2s| {
11199            let p1 = SendMut(o1.as_mut_ptr());
11200            let p2 = SendMut(o2.as_mut_ptr());
11201            let run_range = move |start: usize, end: usize| {
11202                for o in start..end {
11203                    let row = &q[o * cols..(o + 1) * cols];
11204                    let s1 = dot_i8_f32(row, x1s) * row_scale[o];
11205                    let s2 = dot_i8_f32(row, x2s) * row_scale[o];
11206                    // SAFETY: disjoint row ranges per worker.
11207                    unsafe {
11208                        *p1.at(o) = s1;
11209                        *p2.at(o) = s2;
11210                    }
11211                }
11212            };
11213            match pool {
11214                Some(pool) if rows >= 256 => pool.run_rows(rows, &run_range),
11215                _ => run_range(0, rows),
11216            }
11217        });
11218    });
11219}
11220
11221#[derive(Clone, Copy)]
11222struct SendMut(*mut f32);
11223unsafe impl Send for SendMut {}
11224unsafe impl Sync for SendMut {}
11225
11226impl SendMut {
11227    #[inline]
11228    fn at(self, i: usize) -> *mut f32 {
11229        unsafe { self.0.add(i) }
11230    }
11231}
11232
11233#[cfg(test)]
11234mod tests {
11235    /// `q8_round` must be `round().clamp(±127) as i8` bit for bit: every
11236    /// half-integer, their neighbours one ulp either side, the clamp
11237    /// boundary, huge values, infinities and NaN, plus a dense sweep.
11238    #[test]
11239    fn q8_round_is_round_clamp() {
11240        let reference = |t: f32| t.round().clamp(-127.0, 127.0) as i8;
11241        let mut probe = vec![
11242            0.0f32,
11243            -0.0,
11244            f32::NAN,
11245            f32::INFINITY,
11246            f32::NEG_INFINITY,
11247            f32::MAX,
11248            f32::MIN,
11249            1e30,
11250            -1e30,
11251            f32::MIN_POSITIVE,
11252            -f32::MIN_POSITIVE,
11253        ];
11254        for k in -300i32..=300 {
11255            let h = k as f32 * 0.5;
11256            let up = f32::from_bits(h.to_bits() + 1);
11257            let down = f32::from_bits(h.to_bits().wrapping_sub(1));
11258            for t in [h, up, down] {
11259                probe.push(t);
11260                probe.push(-t);
11261            }
11262        }
11263        let mut t = -140.0f32;
11264        while t < 140.0 {
11265            probe.push(t);
11266            t += 0.000_731;
11267        }
11268        for t in probe {
11269            assert_eq!(q8_round(t), reference(t), "t = {t:e} ({:#x})", t.to_bits());
11270        }
11271    }
11272
11273    use super::*;
11274
11275    #[test]
11276    fn q2tp_i8_dot_matches_exact_on_grid() {
11277        // On-grid activations (±1 → sx=1/127, xq=±127 dequantizes
11278        // exactly, no outliers) must make the integer path agree with
11279        // the exact scalar walk to f32 rounding.
11280        let (rows, cols) = (5, 64);
11281        let gpr = cols / GROUP_SIZE;
11282        // Synthetic codes plane + a flat ladder: scales_into is not under
11283        // test here, so drive dot_q2tp_row_i8 / q2tp_row_exact directly
11284        // with hand-made scales.
11285        let chunks: Vec<u8> = (0..rows * gpr * Q2TP_CHUNK)
11286            .map(|i| (i as u32).wrapping_mul(2654435761) as u8)
11287            .collect();
11288        let scales: Vec<f32> = (0..gpr).map(|g| 0.5 + g as f32 * 0.25).collect();
11289        let x: Vec<f32> = (0..cols)
11290            .map(|i| if i % 3 == 0 { -1.0 } else { 1.0 })
11291            .collect();
11292        let act = split_act(&x);
11293        assert!(
11294            act.outliers.is_empty(),
11295            "on-grid input must have no outliers"
11296        );
11297        let gsum = q1_group_sums(&act.xq, gpr);
11298        for r in 0..rows {
11299            let exact = q2tp_row_exact(&chunks, r, gpr, &x, &scales);
11300            let fast = dot_q2tp_row_i8(&chunks, r, gpr, &act.xq, &gsum, &scales) * act.sx;
11301            assert!(
11302                (exact - fast).abs() <= exact.abs() * 1e-5 + 1e-5,
11303                "row {r}: exact {exact} vs i8 {fast}"
11304            );
11305        }
11306    }
11307
11308    #[test]
11309    fn q2tp_affine_fuses_half_scale_correction_without_changing_raw_decode() {
11310        let (rows, cols) = (1usize, GROUP_SIZE);
11311        let mut bytes = vec![0u8; Q2TP_CHUNK + 4 + 1];
11312        // Repeating symbols 0,1,2,0 at unit scale.  q2tp's raw B is
11313        // (c-1.5), while the affine Prism operator is (c-1.0).
11314        bytes[..Q2TP_CHUNK].fill(0x24); // codes 0,1,2,0 in LSB-first order
11315        bytes[Q2TP_CHUNK..Q2TP_CHUNK + 2].copy_from_slice(&0u16.to_le_bytes());
11316        bytes[Q2TP_CHUNK + 2..Q2TP_CHUNK + 4].copy_from_slice(&0u16.to_le_bytes());
11317        bytes[Q2TP_CHUNK + 4] = 1; // dtype16 rung 1 = 1.0
11318        let x = vec![1.0f32; cols];
11319        let mut raw = vec![0.0f32; rows];
11320        let mut affine = vec![0.0f32; rows];
11321        q2tp_matvec_for_test(&bytes, &x, rows, cols, &mut raw);
11322        q2tp_affine_matvec_for_test(&bytes, &x, rows, cols, &mut affine);
11323        assert_eq!(raw, vec![-24.0]);
11324        assert_eq!(affine, vec![-8.0]);
11325        assert!((affine[0] - (raw[0] + 0.5 * cols as f32)).abs() < 1e-6);
11326    }
11327
11328    #[cfg(target_arch = "x86_64")]
11329    #[test]
11330    fn q2tp_avx2_dot_matches_scalar_for_random_patterns() {
11331        // Compare the release AVX2 integer dot against the scalar oracle over
11332        // arbitrary packed bytes/activation signs.  This guards the exact
11333        // table-load path used after rejecting a faster-looking decoder whose
11334        // full-checkpoint greedy output drifted.
11335        if !std::arch::is_x86_feature_detected!("avx2") {
11336            return;
11337        }
11338        let mut seed = 0x9e3779b9u32;
11339        let mut next = || {
11340            seed = seed.wrapping_mul(1664525).wrapping_add(1013904223);
11341            seed
11342        };
11343        for _ in 0..20_000 {
11344            let mut ch = [0u8; Q2TP_CHUNK];
11345            let mut x = [0i8; GROUP_SIZE];
11346            for b in &mut ch {
11347                *b = next() as u8;
11348            }
11349            for v in &mut x {
11350                *v = (next() >> 24) as i8;
11351            }
11352            let mut reference = 0i32;
11353            for (k, &b) in ch.iter().enumerate() {
11354                reference += (b & 3) as i32 * x[k * 4] as i32;
11355                reference += ((b >> 2) & 3) as i32 * x[k * 4 + 1] as i32;
11356                reference += ((b >> 4) & 3) as i32 * x[k * 4 + 2] as i32;
11357                reference += ((b >> 6) & 3) as i32 * x[k * 4 + 3] as i32;
11358            }
11359            // SAFETY: guarded by the runtime AVX2 feature check and fixed
11360            // 8-byte/32-byte slice lengths above.
11361            let got = unsafe { q2tp_code_dot_avx2(&ch, &x) };
11362            assert_eq!(got, reference, "packed q2 lane mismatch");
11363        }
11364    }
11365
11366    #[test]
11367    fn q8_row_dot_fast_matches_scalar() {
11368        // The per-arch fast dot must agree with the exact scalar oracle
11369        // (same contract the fused q8 FFN arm rides on).
11370        let cols = 96;
11371        let row: Vec<u8> = (0..cols)
11372            .map(|i| ((i * 37 % 251) - 125) as i8 as u8)
11373            .collect();
11374        let x: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.13).sin()).collect();
11375        let act = split_act(&x);
11376        let fast = q8_row_dot(&row, &act);
11377        let scalar = q8_row_dot_scalar(&row, &act);
11378        assert!(
11379            (fast - scalar).abs() <= scalar.abs() * 1e-5 + 1e-5,
11380            "fast {fast} vs scalar {scalar}"
11381        );
11382    }
11383
11384    #[test]
11385    fn f32_matvec_matches_matvec_rows_bitexact() {
11386        let (rows, cols) = (300, 40);
11387        let w: Vec<f32> = (0..rows * cols).map(|i| (i as f32 * 0.017).sin()).collect();
11388        let x: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.05).cos()).collect();
11389        let qt = QTensor::from_f32(w.clone(), rows, cols);
11390
11391        let mut a = vec![0.0f32; rows];
11392        matvec_rows(None, &w, &x, &mut a);
11393        let mut b = vec![0.0f32; rows];
11394        qt.matvec(&x, &mut b, None);
11395        assert_eq!(a, b);
11396    }
11397
11398    #[test]
11399    fn sdot_kernel_exact_on_grid() {
11400        // Activations already on the i8 grid (±1 with amax=1 → sx=1/127,
11401        // xq=±127 dequantizes EXACTLY) → the SDOT path must match the
11402        // exact f32 dot to float rounding. This isolates kernel
11403        // correctness from quantization noise.
11404        eprintln!("sdot_enabled = {}", sdot_enabled());
11405        let (rows, cols) = (9, 80); // odd rows → exercises 4-row + tail
11406        let w: Vec<u8> = (0..rows * cols)
11407            .map(|i| (((i * 37) % 251) as i32 - 125) as i8 as u8)
11408            .collect();
11409        let scales: Vec<f32> = (0..rows).map(|o| 0.005 + o as f32 * 0.001).collect();
11410        let x: Vec<f32> = (0..cols)
11411            .map(|i| match i % 3 {
11412                0 => 1.0,
11413                1 => -1.0,
11414                _ => 0.0,
11415            })
11416            .collect();
11417        let mut a = vec![0.0f32; rows];
11418        qmatvec(
11419            &w,
11420            &[],
11421            &scales,
11422            &x,
11423            &[],
11424            TensorDtype::Q8Row,
11425            rows,
11426            cols,
11427            &mut a,
11428            None,
11429        );
11430        for o in 0..rows {
11431            let mut acc = 0.0f32;
11432            for j in 0..cols {
11433                acc += (w[o * cols + j] as i8) as f32 * x[j];
11434            }
11435            let expect = acc * scales[o];
11436            assert!(
11437                (a[o] - expect).abs() < 1e-3 * expect.abs().max(1e-3),
11438                "row {o}: {} vs {expect}",
11439                a[o]
11440            );
11441        }
11442    }
11443
11444    #[test]
11445    fn q1_tbl_fast_path_matches_reference() {
11446        // gpr = 8 exercises the TBL pair-load fast loop, and the LAST
11447        // row's final 4-tile window trips the 4B-overread guard (the
11448        // payload ends exactly at the last tile) — both paths must
11449        // agree with the dequant reference.
11450        let (rows, cols) = (5, 256);
11451        let gpr = cols / GROUP_SIZE;
11452        let mut bytes = Vec::new();
11453        for t in 0..rows * gpr {
11454            let s = 0.007 + (t % 11) as f32 * 0.004;
11455            bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
11456            for j in 0..4 {
11457                bytes.push(((t * 53 + j * 89 + 7) % 249) as u8);
11458            }
11459        }
11460        let x: Vec<f32> = (0..cols)
11461            .map(|i| if (i * 5) % 7 < 3 { 1.0 } else { -1.0 })
11462            .collect();
11463        let mut w = vec![0.0f32; rows * cols];
11464        cortiq_core::quant::dequant_q1(&bytes, &mut w);
11465        let mut got = vec![0.0f32; rows];
11466        q1_matvec(&bytes, &x, rows, cols, &mut got, None);
11467        for o in 0..rows {
11468            let expect: f32 = (0..cols).map(|j| w[o * cols + j] * x[j]).sum();
11469            assert!(
11470                (got[o] - expect).abs() < 1e-3 * expect.abs().max(1e-3),
11471                "row {o}: {} vs {expect}",
11472                got[o]
11473            );
11474        }
11475        // Blocked 1×4 batch (b=5: one quad + remainder) must equal the
11476        // single-matvec path bit-for-bit.
11477        let b = 5usize;
11478        let mut xs_all = Vec::new();
11479        for bi in 0..b {
11480            xs_all.extend(x.iter().map(|v| if bi % 2 == 0 { *v } else { -*v }));
11481        }
11482        let mut mm = vec![0.0f32; b * rows];
11483        q1_matmat(&bytes, &xs_all, b, rows, cols, &mut mm, None);
11484        for bi in 0..b {
11485            let mut single = vec![0.0f32; rows];
11486            q1_matvec(
11487                &bytes,
11488                &xs_all[bi * cols..(bi + 1) * cols],
11489                rows,
11490                cols,
11491                &mut single,
11492                None,
11493            );
11494            assert_eq!(&mm[bi * rows..(bi + 1) * rows], &single[..], "stream {bi}");
11495        }
11496    }
11497
11498    #[test]
11499    fn q1_kernels_match_exact_reference() {
11500        // Synthetic q1 payload: 6-byte tiles [f16 scale][4B bits].
11501        let (rows, cols) = (7, 96);
11502        let gpr = cols / GROUP_SIZE;
11503        let mut bytes = Vec::new();
11504        for t in 0..rows * gpr {
11505            let s = 0.01 + (t % 13) as f32 * 0.003;
11506            bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
11507            for j in 0..4 {
11508                bytes.push(((t * 31 + j * 97) % 251) as u8);
11509            }
11510        }
11511        // On-grid activations (±1, amax 1) → the SDOT path is exact.
11512        let x: Vec<f32> = (0..cols)
11513            .map(|i| if i % 3 == 0 { 1.0 } else { -1.0 })
11514            .collect();
11515        // Reference through the core dequant.
11516        let mut w = vec![0.0f32; rows * cols];
11517        cortiq_core::quant::dequant_q1(&bytes, &mut w);
11518        let mut expect = vec![0.0f32; rows];
11519        for o in 0..rows {
11520            expect[o] = (0..cols).map(|j| w[o * cols + j] * x[j]).sum();
11521        }
11522        let mut got = vec![0.0f32; rows];
11523        q1_matvec(&bytes, &x, rows, cols, &mut got, None);
11524        for o in 0..rows {
11525            assert!(
11526                (got[o] - expect[o]).abs() < 1e-3 * expect[o].abs().max(1e-3),
11527                "row {o}: {} vs {}",
11528                got[o],
11529                expect[o]
11530            );
11531        }
11532        // Pair and batch paths agree with the single path.
11533        let x2: Vec<f32> = x.iter().map(|v| -v).collect();
11534        let (mut a1, mut a2) = (vec![0.0f32; rows], vec![0.0f32; rows]);
11535        q1_matvec2(&bytes, &x, &x2, rows, cols, &mut a1, &mut a2, None);
11536        assert_eq!(a1, got);
11537        let mut xs = x.clone();
11538        xs.extend_from_slice(&x2);
11539        let mut mm = vec![0.0f32; 2 * rows];
11540        q1_matmat(&bytes, &xs, 2, rows, cols, &mut mm, None);
11541        assert_eq!(&mm[..rows], got.as_slice());
11542        assert_eq!(&mm[rows..], a2.as_slice());
11543    }
11544
11545    #[test]
11546    fn repack_is_bit_identical() {
11547        // The interleaved-repack kernel must produce EXACTLY the same
11548        // bits as the mmap-layout kernel: integer accumulation is order-
11549        // exact, the f32 epilogue is identical. Odd rows exercise the
11550        // tail; direct range calls exercise unaligned pool splits.
11551        let (rows, cols) = (267, 96); // 66 groups + 3 tail rows, cols % 16 == 0
11552        let w: Vec<u8> = (0..rows * cols)
11553            .map(|i| (((i * 89) % 253) as i32 - 126) as i8 as u8)
11554            .collect();
11555        let scales: Vec<f32> = (0..rows).map(|o| 0.003 + o as f32 * 0.0007).collect();
11556        let x: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.37).sin() * 2.0).collect();
11557        let rep = q8_repack_layout(&w, rows, cols);
11558        // Group interleave round-trips.
11559        for g in 0..rows / 4 {
11560            for c in 0..cols / 16 {
11561                for lane in 0..4 {
11562                    assert_eq!(
11563                        &rep[g * 4 * cols + c * 64 + lane * 16
11564                            ..g * 4 * cols + c * 64 + lane * 16 + 16],
11565                        &w[(g * 4 + lane) * cols + c * 16..(g * 4 + lane) * cols + c * 16 + 16],
11566                    );
11567                }
11568            }
11569        }
11570        let mut a = vec![0.0f32; rows];
11571        qmatvec(
11572            &w,
11573            &[],
11574            &scales,
11575            &x,
11576            &[],
11577            TensorDtype::Q8Row,
11578            rows,
11579            cols,
11580            &mut a,
11581            None,
11582        );
11583        let mut b = vec![0.0f32; rows];
11584        qmatvec(
11585            &w,
11586            &rep,
11587            &scales,
11588            &x,
11589            &[],
11590            TensorDtype::Q8Row,
11591            rows,
11592            cols,
11593            &mut b,
11594            None,
11595        );
11596        assert_eq!(a, b, "full-range repack output diverged");
11597
11598        #[cfg(target_arch = "aarch64")]
11599        if sdot_enabled() {
11600            // Unaligned range split (pool workers get arbitrary bounds).
11601            let act = split_act(&x);
11602            let mut c1 = vec![0.0f32; rows];
11603            let mut c2 = vec![0.0f32; rows];
11604            q8_range_sdot(
11605                &w,
11606                &[],
11607                &scales,
11608                &act,
11609                cols,
11610                SendMut(c1.as_mut_ptr()),
11611                3,
11612                rows - 2,
11613            );
11614            q8_range_sdot(
11615                &w,
11616                &rep,
11617                &scales,
11618                &act,
11619                cols,
11620                SendMut(c2.as_mut_ptr()),
11621                3,
11622                rows - 2,
11623            );
11624            assert_eq!(c1, c2, "unaligned-range repack output diverged");
11625        }
11626    }
11627
11628    #[test]
11629    fn sdot_a8w8_noise_is_bounded() {
11630        // Off-grid activations: A8 quantization noise must stay small in
11631        // relative L2 over the whole output (realistic accuracy contract;
11632        // vmfcore measured argmax-identical decode on real models).
11633        let (rows, cols) = (16, 512);
11634        let w: Vec<u8> = (0..rows * cols)
11635            .map(|i| (((i * 37) % 251) as i32 - 125) as i8 as u8)
11636            .collect();
11637        let scales = vec![0.01f32; rows];
11638        let x: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.21).sin()).collect();
11639        let mut a = vec![0.0f32; rows];
11640        qmatvec(
11641            &w,
11642            &[],
11643            &scales,
11644            &x,
11645            &[],
11646            TensorDtype::Q8Row,
11647            rows,
11648            cols,
11649            &mut a,
11650            None,
11651        );
11652        let (mut num, mut den) = (0f64, 0f64);
11653        for o in 0..rows {
11654            let mut acc = 0.0f32;
11655            for j in 0..cols {
11656                acc += (w[o * cols + j] as i8) as f32 * x[j];
11657            }
11658            let expect = acc * scales[o];
11659            num += ((a[o] - expect) as f64).powi(2);
11660            den += (expect as f64).powi(2);
11661        }
11662        let rel = (num / den.max(1e-12)).sqrt();
11663        assert!(rel < 0.05, "A8W8 relative L2 error too high: {rel}");
11664    }
11665
11666    #[test]
11667    fn i8_dot_neon_matches_scalar() {
11668        let n = 100;
11669        let w: Vec<u8> = (0..n).map(|i| ((i * 37 + 11) % 251) as u8).collect();
11670        let x: Vec<f32> = (0..n).map(|i| (i as f32 * 0.13).sin()).collect();
11671        let mut scalar = 0.0f32;
11672        for j in 0..n {
11673            scalar += (w[j] as i8) as f32 * x[j];
11674        }
11675        let fast = dot_i8_f32(&w, &x);
11676        assert!((scalar - fast).abs() < 1e-3 * scalar.abs().max(1.0));
11677    }
11678
11679    /// Fused vbit matvec must match full dequant_vbit + dense matvec.
11680    #[test]
11681    fn vbitmatvec_matches_full_dequant() {
11682        let (rows, cols) = (6, 64);
11683        let ng = cols / GROUP_SIZE;
11684        // Hand-craft: bits per row, f16 scales, packed rows.
11685        let bits: Vec<u8> = vec![3, 4, 5, 6, 8, 4];
11686        let mut bytes = bits.clone();
11687        for g in 0..rows * ng {
11688            let s = 0.02 + 0.001 * g as f32;
11689            bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
11690        }
11691        for r in 0..rows {
11692            let b = bits[r] as usize;
11693            let (mut acc, mut nb) = (0u64, 0usize);
11694            let mut rowbytes = Vec::new();
11695            for i in 0..cols {
11696                let v = ((i * 7 + r * 13) % (1 << b)) as u64;
11697                acc = (acc << b) | v;
11698                nb += b;
11699                while nb >= 8 {
11700                    nb -= 8;
11701                    rowbytes.push(((acc >> nb) & 0xFF) as u8);
11702                }
11703            }
11704            if nb > 0 {
11705                rowbytes.push(((acc << (8 - nb)) & 0xFF) as u8);
11706            }
11707            bytes.extend_from_slice(&rowbytes);
11708        }
11709        let x: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.19).sin()).collect();
11710
11711        let mut reference = vec![0f32; rows * cols];
11712        cortiq_core::quant::dequant_vbit(&bytes, rows, cols, &mut reference).unwrap();
11713        let mut expect = vec![0f32; rows];
11714        for r in 0..rows {
11715            expect[r] = reference[r * cols..(r + 1) * cols]
11716                .iter()
11717                .zip(&x)
11718                .map(|(w, xv)| w * xv)
11719                .sum();
11720        }
11721        let mut got = vec![0f32; rows];
11722        let offsets = vbit_row_offsets(&bytes, rows, cols);
11723        vbitmatvec(&bytes, &offsets, &x, rows, cols, &mut got, None);
11724        // SDOT path quantizes activations to i8 (A8W8): bounded noise,
11725        // same contract as q8 (exact path is pinned by CMF_SDOT=0 in
11726        // the golden-parity gate).
11727        let tol = if a8w8_enabled() { 6e-2 } else { 1e-4 };
11728        let scale = expect.iter().fold(0f32, |m, v| m.max(v.abs())).max(1e-6);
11729        for r in 0..rows {
11730            assert!(
11731                (got[r] - expect[r]).abs() < tol * scale,
11732                "row {r}: {} vs {}",
11733                got[r],
11734                expect[r]
11735            );
11736        }
11737    }
11738
11739    /// Fused q4 matvec must match the reference full-dequant + dense
11740    /// matvec bit-for-bit in structure (same f32 math, group order).
11741    /// vbit matmat: the blocked 1×4 leg must match the per-row path
11742    /// (paired env toggle; larger shape so both code paths engage).
11743    #[test]
11744    #[cfg(target_arch = "x86_64")]
11745    fn vbit_matmat_blocked_matches_per_row() {
11746        let (rows, cols, b) = (64usize, 128usize, 9usize);
11747        let ng = cols / GROUP_SIZE;
11748        let bits: Vec<u8> = (0..rows).map(|r| [3u8, 4, 5, 6][r % 4]).collect();
11749        let mut bytes = bits.clone();
11750        for g in 0..rows * ng {
11751            let sc = 0.02 + 0.0005 * g as f32;
11752            bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(sc).to_le_bytes());
11753        }
11754        for r in 0..rows {
11755            let bw = bits[r] as usize;
11756            let (mut acc, mut nb) = (0u64, 0usize);
11757            let mut rowbytes = Vec::new();
11758            for i in 0..cols {
11759                let v = ((i * 7 + r * 13) % (1 << bw)) as u64;
11760                acc = (acc << bw) | v;
11761                nb += bw;
11762                while nb >= 8 {
11763                    nb -= 8;
11764                    rowbytes.push(((acc >> nb) & 0xFF) as u8);
11765                }
11766            }
11767            if nb > 0 {
11768                rowbytes.push(((acc << (8 - nb)) & 0xFF) as u8);
11769            }
11770            bytes.extend_from_slice(&rowbytes);
11771        }
11772        let x: Vec<f32> = (0..b * cols)
11773            .map(|i| ((i * 13 + 7) % 97) as f32 / 97.0 - 0.5)
11774            .collect();
11775        let offsets = vbit_row_offsets(&bytes, rows, cols);
11776        let mut y_a = vec![0f32; b * rows];
11777        let mut y_b = vec![0f32; b * rows];
11778        unsafe { std::env::set_var("CMF_X86_BLOCKED", "1") };
11779        vbitmatmat(&bytes, &offsets, &x, b, rows, cols, &mut y_a, None);
11780        unsafe { std::env::set_var("CMF_X86_BLOCKED", "0") };
11781        vbitmatmat(&bytes, &offsets, &x, b, rows, cols, &mut y_b, None);
11782        unsafe { std::env::remove_var("CMF_X86_BLOCKED") };
11783        let max_d = y_a
11784            .iter()
11785            .zip(&y_b)
11786            .map(|(p, q)| (p - q).abs())
11787            .fold(0.0f32, f32::max);
11788        assert!(max_d < 1e-4, "vbit blocked ≠ per-row: max|Δ| = {max_d}");
11789    }
11790
11791    /// q4t blocked 1×4 (SDOT on ARM, AVX2 on x86) must equal the
11792    /// per-row path exactly: same nibble unpack, same group order,
11793    /// same f32 accumulation — batch == matvec bit-for-bit. b=9 covers
11794    /// two full 1×4 blocks plus a remainder through the single-row
11795    /// kernel. (Both paths produce identical output, so the shared
11796    /// CMF_X86_BLOCKED env var racing with other tests cannot flip
11797    /// the verdict — worst case both sides take the same path.)
11798    #[test]
11799    fn q4t_matmat_blocked_matches_per_row() {
11800        let (rows, cols, b) = (16usize, 64usize, 9usize);
11801        let gpr = cols / GROUP_SIZE;
11802        let mut bytes = vec![0u8; rows * gpr * Q4_TILE];
11803        for r in 0..rows {
11804            for g in 0..gpr {
11805                let t = (r * gpr + g) * Q4_TILE;
11806                let sc = 0.02 + 0.001 * (r * gpr + g) as f32;
11807                bytes[t..t + 2].copy_from_slice(&cortiq_core::quant::f32_to_f16(sc).to_le_bytes());
11808                for k in 0..16 {
11809                    bytes[t + 2 + k] = ((r * 31 + g * 7 + k * 13) % 251) as u8;
11810                }
11811            }
11812        }
11813        let x: Vec<f32> = (0..b * cols)
11814            .map(|i| ((i * 13 + 7) % 97) as f32 / 97.0 - 0.5)
11815            .collect();
11816        let mut y_blk = vec![0f32; b * rows];
11817        let mut y_row = vec![0f32; b * rows];
11818        unsafe { std::env::set_var("CMF_X86_BLOCKED", "1") };
11819        q4t_matmat(&bytes, &x, b, rows, cols, &mut y_blk, None);
11820        unsafe { std::env::set_var("CMF_X86_BLOCKED", "0") };
11821        q4t_matmat(&bytes, &x, b, rows, cols, &mut y_row, None);
11822        unsafe { std::env::remove_var("CMF_X86_BLOCKED") };
11823        assert_eq!(y_blk, y_row, "q4t blocked 1x4 ≠ per-row");
11824    }
11825
11826    /// The wide-batch Accelerate arm of q4t_matmat vs a brute-force
11827    /// f32 dequant matmul: both are f32 GEMMs, so only reduction
11828    /// order differs — tight tolerance.
11829    /// A synthetic q4tp payload: random nibbles plus a per-row ladder whose
11830    /// span varies row to row, so the codes actually exercise the full 0..31
11831    /// range rather than clustering on one rung.
11832    fn synth_q4tp(rows: usize, cols: usize) -> Vec<u8> {
11833        use cortiq_core::quant::{f32_to_f16, q4tp_code_stride, q4tp_put_code};
11834        let gpr = cols / GROUP_SIZE;
11835        let stride = q4tp_code_stride(gpr);
11836        let (params_off, codes_off, _) = q4tp_sections(rows, cols);
11837        let mut b = vec![0u8; codes_off + rows * stride];
11838        for r in 0..rows {
11839            for g in 0..gpr {
11840                let t = (r * gpr + g) * Q4TP_NIB;
11841                for k in 0..16 {
11842                    b[t + k] = ((r * 31 + g * 7 + k * 13) % 251) as u8;
11843                }
11844            }
11845            let lo = -6.0 - 0.03 * (r % 17) as f32;
11846            let step = 0.01 + 0.004 * (r % 11) as f32;
11847            let p = params_off + r * 4;
11848            b[p..p + 2].copy_from_slice(&f32_to_f16(lo).to_le_bytes());
11849            b[p + 2..p + 4].copy_from_slice(&f32_to_f16(step).to_le_bytes());
11850            let crow = &mut b[codes_off + r * stride..codes_off + (r + 1) * stride];
11851            for g in 0..gpr {
11852                q4tp_put_code(crow, g, (r * 5 + g * 3) % 32);
11853            }
11854        }
11855        b
11856    }
11857
11858    /// The same weights re-expressed as q4_tiled, so the proven kernel can
11859    /// be the reference: each tile stores the ladder scale its code selects.
11860    /// Only the f16 rounding of that scale separates the two payloads.
11861    fn q4tp_as_q4t(bytes: &[u8], rows: usize, cols: usize) -> Vec<u8> {
11862        let gpr = cols / GROUP_SIZE;
11863        let v = Q4tpView::new(bytes, rows, cols);
11864        let mut out = vec![0u8; rows * gpr * Q4_TILE];
11865        let mut sc = vec![0f32; gpr];
11866        for r in 0..rows {
11867            v.scales_into(r, gpr, &mut sc);
11868            for g in 0..gpr {
11869                let t = (r * gpr + g) * Q4_TILE;
11870                let s = sc[g];
11871                out[t..t + 2].copy_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
11872                let src = (r * gpr + g) * Q4TP_NIB;
11873                out[t + 2..t + Q4_TILE].copy_from_slice(&v.nib[src..src + Q4TP_NIB]);
11874            }
11875        }
11876        out
11877    }
11878
11879    /// The exact (`CMF_SDOT=0`) path must reproduce `dequant_q4tp` to f32
11880    /// rounding — that scalar routine is the format's definition, and the
11881    /// kernels re-derive the scale from the ladder independently. Call the
11882    /// row kernel directly: `matmat` picks the int8 arm when a8w8 is on,
11883    /// so routing through it would test the other path by accident.
11884    #[test]
11885    fn q4tp_exact_path_matches_dequant_reference() {
11886        let (rows, cols) = (256usize, 512usize);
11887        let gpr = cols / GROUP_SIZE;
11888        let bytes = synth_q4tp(rows, cols);
11889        let mut w = vec![0f32; rows * cols];
11890        cortiq_core::quant::dequant_q4tp(&bytes, rows, cols, &mut w);
11891
11892        let x: Vec<f32> = (0..cols)
11893            .map(|i| ((i * 13 + 7) % 97) as f32 / 97.0 - 0.5)
11894            .collect();
11895        let v = Q4tpView::new(&bytes, rows, cols);
11896        let mut sc = vec![0f32; gpr];
11897        for r in 0..rows {
11898            v.scales_into(r, gpr, &mut sc);
11899            let got = q4tp_row_exact(v.nib, r, gpr, &x, &sc);
11900            let want: f32 = (0..cols).map(|c| w[r * cols + c] * x[c]).sum();
11901            // These dot products cancel down to ~1e-3 from terms of ~5e-2, so
11902            // the meaningful yardstick is the summed magnitude, not the result:
11903            // against the result any reordering of a 512-term f32 sum "fails".
11904            let mag: f32 = (0..cols).map(|c| (w[r * cols + c] * x[c]).abs()).sum();
11905            assert!(
11906                (got - want).abs() <= 1e-5 * mag,
11907                "row {r}: kernel {got} vs dequant {want}"
11908            );
11909        }
11910    }
11911
11912    /// The int8 (a8w8) path can't be checked against an f32 reference — the
11913    /// activation quantization dominates. Check it against the q4t kernel it
11914    /// was ported from instead, on payloads holding the same weights: that
11915    /// isolates exactly what the port could break (16 B stride, ladder
11916    /// lookup, nibble unpack) from what it deliberately shares.
11917    #[test]
11918    fn q4tp_matvec_matches_the_q4t_kernel_it_was_ported_from() {
11919        let (rows, cols) = (256usize, 512usize);
11920        let bytes = synth_q4tp(rows, cols);
11921        let twin = q4tp_as_q4t(&bytes, rows, cols);
11922        let x: Vec<f32> = (0..cols)
11923            .map(|i| ((i * 13 + 7) % 97) as f32 / 97.0 - 0.5)
11924            .collect();
11925
11926        let mut got = vec![0f32; rows];
11927        q4tp_matvec(&bytes, &x, rows, cols, &mut got, None);
11928        let mut want = vec![0f32; rows];
11929        q4t_matvec(&twin, &x, rows, cols, &mut want, None);
11930
11931        // Scale is f16 in the twin and f32 here, so allow that rounding on
11932        // top of the summed magnitude (same cancellation argument as above).
11933        let mut w = vec![0f32; rows * cols];
11934        cortiq_core::quant::dequant_q4tp(&bytes, rows, cols, &mut w);
11935        for r in 0..rows {
11936            let mag: f32 = (0..cols).map(|c| (w[r * cols + c] * x[c]).abs()).sum();
11937            assert!(
11938                (got[r] - want[r]).abs() <= 1e-3 * mag,
11939                "row {r}: q4tp {} vs q4t {}",
11940                got[r],
11941                want[r]
11942            );
11943        }
11944    }
11945
11946    /// `matmat` carries three arms (Accelerate, blocked int8 1x4, scalar).
11947    /// Batch 5 crosses the blocked kernel's stride, so this exercises the
11948    /// 1x4 path AND its scalar tail in one run — the blocked kernel is new
11949    /// code and its four accumulators are exactly what tends to go wrong.
11950    #[test]
11951    fn q4tp_matmat_matches_the_q4t_kernel_it_was_ported_from() {
11952        let (rows, cols, b) = (256usize, 512usize, 5usize);
11953        let bytes = synth_q4tp(rows, cols);
11954        let twin = q4tp_as_q4t(&bytes, rows, cols);
11955        let xs: Vec<f32> = (0..b * cols)
11956            .map(|i| ((i * 29 + 11) % 89) as f32 / 89.0 - 0.5)
11957            .collect();
11958
11959        let mut got = vec![0f32; b * rows];
11960        q4tp_matmat(&bytes, &xs, b, rows, cols, &mut got, None);
11961        let mut want = vec![0f32; b * rows];
11962        q4t_matmat(&twin, &xs, b, rows, cols, &mut want, None);
11963
11964        let mut w = vec![0f32; rows * cols];
11965        cortiq_core::quant::dequant_q4tp(&bytes, rows, cols, &mut w);
11966        for t in 0..b {
11967            for r in 0..rows {
11968                let mag: f32 = (0..cols)
11969                    .map(|c| (w[r * cols + c] * xs[t * cols + c]).abs())
11970                    .sum();
11971                let (g, wa) = (got[t * rows + r], want[t * rows + r]);
11972                assert!(
11973                    (g - wa).abs() <= 1e-3 * mag,
11974                    "batch {t} row {r}: q4tp {g} vs q4t {wa}"
11975                );
11976            }
11977        }
11978    }
11979
11980    #[test]
11981    fn q4tp_matvec2_matches_the_single_stream_kernel() {
11982        let (rows, cols) = (128usize, 256usize);
11983        let gpr = cols / GROUP_SIZE;
11984        let bytes = synth_q4tp(rows, cols);
11985        let xs: Vec<f32> = (0..2 * cols)
11986            .map(|i| ((i * 29 + 11) % 89) as f32 / 89.0 - 0.5)
11987            .collect();
11988
11989        let (mut o1, mut o2) = (vec![0f32; rows], vec![0f32; rows]);
11990        q4tp_matvec2(
11991            &bytes,
11992            &xs[..cols],
11993            &xs[cols..],
11994            rows,
11995            cols,
11996            &mut o1,
11997            &mut o2,
11998            None,
11999        );
12000
12001        // matvec2 takes the exact path for both streams, so the single-row
12002        // kernel is an exact reference — no tolerance for path differences.
12003        let v = Q4tpView::new(&bytes, rows, cols);
12004        let mut sc = vec![0f32; gpr];
12005        for r in 0..rows {
12006            v.scales_into(r, gpr, &mut sc);
12007            assert_eq!(o1[r], q4tp_row_exact(v.nib, r, gpr, &xs[..cols], &sc));
12008            assert_eq!(o2[r], q4tp_row_exact(v.nib, r, gpr, &xs[cols..], &sc));
12009        }
12010    }
12011
12012    /// q4tp must not COST speed — it exists to save bytes, and a format that
12013    /// trades 7% of a file for a slower model is a bad trade. This guard is
12014    /// here because correctness tests happily passed while `q4tp_matmat` was
12015    /// missing its int8 and Accelerate arms and the model ran 5x slower.
12016    /// Measured on M-series: 0.97-1.04x, i.e. parity (16 B tiles are better
12017    /// aligned than q4t's 18 B, which pays for the scale indirection).
12018    #[test]
12019    fn q4tp_matvec_keeps_pace_with_q4t() {
12020        let (rows, cols) = (4096usize, 3072usize);
12021        let bytes = synth_q4tp(rows, cols);
12022        let twin = q4tp_as_q4t(&bytes, rows, cols);
12023        let x: Vec<f32> = (0..cols).map(|i| (i % 97) as f32 / 97.0 - 0.5).collect();
12024        let mut o = vec![0f32; rows];
12025        let n = 12;
12026        let mut best = (f64::MAX, f64::MAX);
12027        // Interleaved A/B, minimum statistic: this machine throttles, and a
12028        // mean over a thermal ramp reliably indicts whichever ran second.
12029        for _ in 0..3 {
12030            let t0 = std::time::Instant::now();
12031            for _ in 0..n {
12032                q4t_matvec(&twin, &x, rows, cols, &mut o, None);
12033            }
12034            best.0 = best.0.min(t0.elapsed().as_secs_f64());
12035            let t0 = std::time::Instant::now();
12036            for _ in 0..n {
12037                q4tp_matvec(&bytes, &x, rows, cols, &mut o, None);
12038            }
12039            best.1 = best.1.min(t0.elapsed().as_secs_f64());
12040        }
12041        let ratio = best.1 / best.0;
12042        println!(
12043            "q4t {:.3} ms | q4tp {:.3} ms | {ratio:.2}x",
12044            best.0 * 1e3 / n as f64,
12045            best.1 * 1e3 / n as f64
12046        );
12047        assert!(ratio < 2.0, "q4tp matvec {ratio:.2}x slower than q4t");
12048    }
12049
12050    #[cfg(target_os = "macos")]
12051    #[test]
12052    fn q4t_matmat_accel_matches_dequant_reference() {
12053        if !accel_gemm_enabled() {
12054            return; // CMF_ACCEL=0
12055        }
12056        let (rows, cols, b) = (512usize, 1024usize, 8usize); // ≥500K → accel arm
12057        let gpr = cols / GROUP_SIZE;
12058        let mut bytes = vec![0u8; rows * gpr * Q4_TILE];
12059        for r in 0..rows {
12060            for g in 0..gpr {
12061                let t = (r * gpr + g) * Q4_TILE;
12062                let sc = 0.02 + 0.0005 * ((r * gpr + g) % 64) as f32;
12063                bytes[t..t + 2].copy_from_slice(&cortiq_core::quant::f32_to_f16(sc).to_le_bytes());
12064                for k in 0..16 {
12065                    bytes[t + 2 + k] = ((r * 31 + g * 7 + k * 13) % 251) as u8;
12066                }
12067            }
12068        }
12069        let x: Vec<f32> = (0..b * cols)
12070            .map(|i| ((i * 13 + 7) % 97) as f32 / 97.0 - 0.5)
12071            .collect();
12072        let mut got = vec![0f32; b * rows];
12073        q4t_matmat(&bytes, &x, b, rows, cols, &mut got, None);
12074        // Brute-force reference off the same tiles.
12075        let mut w = vec![0f32; rows * cols];
12076        for r in 0..rows {
12077            for g in 0..gpr {
12078                let t = (r * gpr + g) * Q4_TILE;
12079                let s = f16_to_f32(u16::from_le_bytes([bytes[t], bytes[t + 1]]));
12080                for (k, &bb) in bytes[t + 2..t + Q4_TILE].iter().enumerate() {
12081                    w[r * cols + g * GROUP_SIZE + k * 2] = ((bb & 0x0F) as f32 - 8.0) * s;
12082                    w[r * cols + g * GROUP_SIZE + k * 2 + 1] =
12083                        (((bb >> 4) & 0x0F) as f32 - 8.0) * s;
12084                }
12085            }
12086        }
12087        for bi in 0..b {
12088            for r in 0..rows {
12089                let want: f32 = (0..cols).map(|j| x[bi * cols + j] * w[r * cols + j]).sum();
12090                let d = (got[bi * rows + r] - want).abs();
12091                assert!(
12092                    d <= want.abs().max(1.0) * 1e-4,
12093                    "accel q4t GEMM diverged at ({bi},{r}): {} vs {want}",
12094                    got[bi * rows + r]
12095                );
12096            }
12097        }
12098    }
12099
12100    #[test]
12101    fn q4matvec_matches_full_dequant() {
12102        let (rows, cols) = (8, 64);
12103        let groups = rows * cols / GROUP_SIZE;
12104        // Hand-craft a q4_block blob: nibbles then f16 scales.
12105        let mut bytes = Vec::with_capacity(groups * 16 + groups * 2);
12106        for i in 0..groups * 16 {
12107            bytes.push((((i * 7 + 3) % 256) & 0xFF) as u8);
12108        }
12109        for g in 0..groups {
12110            let s = 0.01 + 0.003 * g as f32;
12111            bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
12112        }
12113        let x: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.17).sin()).collect();
12114
12115        let mut reference = vec![0.0f32; rows * cols];
12116        cortiq_core::quant::dequant_q4_block(&bytes, &mut reference);
12117        let mut expect = vec![0.0f32; rows];
12118        for r in 0..rows {
12119            expect[r] = reference[r * cols..(r + 1) * cols]
12120                .iter()
12121                .zip(&x)
12122                .map(|(w, xv)| w * xv)
12123                .sum();
12124        }
12125
12126        let mut got = vec![0.0f32; rows];
12127        q4matvec(&bytes, &x, rows, cols, &mut got, None);
12128        // SDOT path quantizes activations to i8 (A8W8): bounded noise,
12129        // same contract as q8/vbit (exact path is pinned by CMF_SDOT=0
12130        // in the golden-parity gate).
12131        let tol = if a8w8_enabled() { 6e-2 } else { 1e-4 };
12132        let scale = expect.iter().fold(0f32, |m, v| m.max(v.abs())).max(1.0);
12133        for r in 0..rows {
12134            assert!(
12135                (got[r] - expect[r]).abs() < tol * scale,
12136                "row {r}: {} vs {}",
12137                got[r],
12138                expect[r]
12139            );
12140        }
12141    }
12142
12143    /// Fused two-input vbit matvec must equal two single matvecs exactly
12144    /// (same per-lane accumulation order on both scalar and SDOT paths).
12145    #[test]
12146    fn vbitmatvec2_equals_two_singles() {
12147        let (rows, cols) = (6, 64);
12148        let ng = cols / GROUP_SIZE;
12149        let bits: Vec<u8> = vec![3, 4, 5, 6, 8, 4];
12150        let mut bytes = bits.clone();
12151        for g in 0..rows * ng {
12152            let s = 0.02 + 0.001 * g as f32;
12153            bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
12154        }
12155        for r in 0..rows {
12156            let b = bits[r] as usize;
12157            let (mut acc, mut nb) = (0u64, 0usize);
12158            let mut rowbytes = Vec::new();
12159            for i in 0..cols {
12160                let v = ((i * 7 + r * 13) % (1 << b)) as u64;
12161                acc = (acc << b) | v;
12162                nb += b;
12163                while nb >= 8 {
12164                    nb -= 8;
12165                    rowbytes.push(((acc >> nb) & 0xFF) as u8);
12166                }
12167            }
12168            if nb > 0 {
12169                rowbytes.push(((acc << (8 - nb)) & 0xFF) as u8);
12170            }
12171            bytes.extend_from_slice(&rowbytes);
12172        }
12173        let x1: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.19).sin()).collect();
12174        let x2: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.11).cos()).collect();
12175        let offsets = vbit_row_offsets(&bytes, rows, cols);
12176
12177        let (mut a1, mut a2) = (vec![0f32; rows], vec![0f32; rows]);
12178        vbitmatvec(&bytes, &offsets, &x1, rows, cols, &mut a1, None);
12179        vbitmatvec(&bytes, &offsets, &x2, rows, cols, &mut a2, None);
12180        let (mut b1, mut b2) = (vec![0f32; rows], vec![0f32; rows]);
12181        vbitmatvec2(
12182            &bytes, &offsets, &x1, &x2, rows, cols, &mut b1, &mut b2, None,
12183        );
12184        assert_eq!(a1, b1, "fused vbit lane 1 must be bit-identical");
12185        assert_eq!(a2, b2, "fused vbit lane 2 must be bit-identical");
12186    }
12187
12188    /// Fused two-input q4 matvec must equal two single matvecs exactly.
12189    #[test]
12190    fn q4matvec2_equals_two_singles() {
12191        let (rows, cols) = (8, 128);
12192        let groups = rows * cols / GROUP_SIZE;
12193        let mut bytes = Vec::with_capacity(groups * 16 + groups * 2);
12194        for i in 0..groups * 16 {
12195            bytes.push((((i * 7 + 3) % 256) & 0xFF) as u8);
12196        }
12197        for g in 0..groups {
12198            let s = 0.01 + 0.003 * g as f32;
12199            bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
12200        }
12201        // Include an outlier channel so the SDOT correction path is
12202        // exercised in the pair kernel too.
12203        let mut x1: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.17).sin()).collect();
12204        x1[9] = 250.0;
12205        let x2: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.23).cos()).collect();
12206
12207        let (mut a1, mut a2) = (vec![0f32; rows], vec![0f32; rows]);
12208        q4matvec(&bytes, &x1, rows, cols, &mut a1, None);
12209        q4matvec(&bytes, &x2, rows, cols, &mut a2, None);
12210        let (mut b1, mut b2) = (vec![0f32; rows], vec![0f32; rows]);
12211        q4matvec2(&bytes, &x1, &x2, rows, cols, &mut b1, &mut b2, None);
12212        assert_eq!(a1, b1, "fused q4 lane 1 must be bit-identical");
12213        assert_eq!(a2, b2, "fused q4 lane 2 must be bit-identical");
12214    }
12215
12216    /// Multi-matrix job must equal separate matvecs exactly — same
12217    /// kernels, only the dispatch is fused.
12218    #[test]
12219    fn matvec_many_equals_separate_matvecs() {
12220        use crate::pool::Pool;
12221        let (r1, r2, cols) = (300, 200, 64);
12222        let mk = |salt: usize, rows: usize| {
12223            QTensor::from_f32(
12224                (0..rows * cols)
12225                    .map(|i| ((i * 7 + salt) % 97) as f32 / 97.0 - 0.5)
12226                    .collect(),
12227                rows,
12228                cols,
12229            )
12230        };
12231        let (a, b) = (mk(1, r1), mk(5, r2));
12232        let x: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.11).sin()).collect();
12233        let pool = Pool::new(3);
12234
12235        let (mut ea, mut eb) = (vec![0f32; r1], vec![0f32; r2]);
12236        a.matvec(&x, &mut ea, Some(&pool));
12237        b.matvec(&x, &mut eb, Some(&pool));
12238        let (mut ga, mut gb) = (vec![0f32; r1], vec![0f32; r2]);
12239        QTensor::matvec_many([&a, &b], &x, [&mut ga, &mut gb], Some(&pool));
12240        assert_eq!(ea, ga, "fused multi-matrix lane 1 must be bit-identical");
12241        assert_eq!(eb, gb, "fused multi-matrix lane 2 must be bit-identical");
12242    }
12243
12244    /// The public Q4TP operator must take the real mapped matvec_many arm,
12245    /// rather than the F32 fallback above.  Build a tiny valid CMF so both
12246    /// handles retain their mmap payloads, then compare the fused dispatch
12247    /// with two ordinary mapped matvec calls bit-for-bit.
12248    #[test]
12249    fn q4tp_matvec_many_equals_separate_matvecs() {
12250        use crate::pool::Pool;
12251        use cortiq_core::{CMF_VERSION, CmfHeader, CmfModel, QuantType, TensorSpec};
12252
12253        let (r1, r2, cols) = (300usize, 200usize, 64usize);
12254        let arch: cortiq_core::ModelArch = serde_json::from_value(serde_json::json!({
12255            "arch_name": "tiny-q4tp",
12256            "hidden_size": cols,
12257            "intermediate_size": cols * 2,
12258            "num_layers": 1,
12259            "num_attention_heads": 2,
12260            "num_kv_heads": 1,
12261            "head_dim": 32,
12262            "vocab_size": r1,
12263            "layer_types": ["FullAttention"],
12264            "rms_norm_eps": 1e-6,
12265            "max_position_embeddings": 8,
12266            "linear_conv_kernel_dim": 0,
12267            "linear_num_key_heads": 0,
12268            "linear_num_value_heads": 0
12269        }))
12270        .unwrap();
12271        let header = CmfHeader {
12272            format: "cmf".into(),
12273            version: CMF_VERSION,
12274            arch,
12275            quant_type: QuantType::Q4Block,
12276            provenance: None,
12277            tokenizer_config: None,
12278            section_hashes: None,
12279            skills: Vec::new(),
12280            shard: None,
12281            calibration: None,
12282            routing: None,
12283        };
12284        let specs = [
12285            TensorSpec {
12286                name: "q".into(),
12287                dtype: TensorDtype::Q4TiledP,
12288                shape: vec![r1, cols],
12289                data: synth_q4tp(r1, cols),
12290            },
12291            TensorSpec {
12292                name: "kv".into(),
12293                dtype: TensorDtype::Q4TiledP,
12294                shape: vec![r2, cols],
12295                data: synth_q4tp(r2, cols),
12296            },
12297        ];
12298        let dir = std::env::temp_dir().join(format!("cmf-q4tp-many-{}", std::process::id()));
12299        std::fs::create_dir_all(&dir).unwrap();
12300        let path = dir.join("m.cmf");
12301        CmfModel::write(&path, &header, &specs, None, None).unwrap();
12302        let model = Arc::new(CmfModel::open(&path).unwrap());
12303        let (a, b) = (
12304            QTensor::from_model(&model, "q").unwrap(),
12305            QTensor::from_model(&model, "kv").unwrap(),
12306        );
12307        assert_eq!(a.model_dtype(), Some(TensorDtype::Q4TiledP));
12308        assert_eq!(b.model_dtype(), Some(TensorDtype::Q4TiledP));
12309        let x: Vec<f32> = (0..cols)
12310            .map(|i| ((i * 17 + 3) % 97) as f32 / 97.0 - 0.5)
12311            .collect();
12312        let pool = Pool::new(3);
12313        let (mut ea, mut eb) = (vec![0.0f32; r1], vec![0.0f32; r2]);
12314        a.matvec(&x, &mut ea, Some(&pool));
12315        b.matvec(&x, &mut eb, Some(&pool));
12316        let (mut ga, mut gb) = (vec![0.0f32; r1], vec![0.0f32; r2]);
12317        QTensor::matvec_many([&a, &b], &x, [&mut ga, &mut gb], Some(&pool));
12318        assert_eq!(ea, ga, "Q4TP fused lane 1 must be bit-identical");
12319        assert_eq!(eb, gb, "Q4TP fused lane 2 must be bit-identical");
12320        let _ = std::fs::remove_dir_all(&dir);
12321    }
12322
12323    /// The MiMo speculative verify's kernels: several tokens' MoE through
12324    /// `moe_gate_up_rows` / `moe_down_rows` (+ the caller's route-order sum)
12325    /// is bit-identical to each token alone through `moe_gate_up_many` /
12326    /// `moe_down_many` (decode), and a row-exact `q4tp_matmat` of five
12327    /// tokens (wide enough for the blocked tiles) equals five matvecs.
12328    #[test]
12329    fn multi_token_moe_rows_equal_single_token_decode() {
12330        use crate::pool::Pool;
12331        use cortiq_core::{CMF_VERSION, CmfHeader, CmfModel, QuantType, TensorSpec};
12332
12333        let (h, inter, ne) = (64usize, 128usize, 3usize);
12334        let arch: cortiq_core::ModelArch = serde_json::from_value(serde_json::json!({
12335            "arch_name": "tiny-q4tp-moe",
12336            "hidden_size": h,
12337            "intermediate_size": inter,
12338            "num_layers": 1,
12339            "num_attention_heads": 2,
12340            "num_kv_heads": 1,
12341            "head_dim": 32,
12342            "vocab_size": 8,
12343            "layer_types": ["FullAttention"],
12344            "rms_norm_eps": 1e-6,
12345            "max_position_embeddings": 8,
12346            "linear_conv_kernel_dim": 0,
12347            "linear_num_key_heads": 0,
12348            "linear_num_value_heads": 0
12349        }))
12350        .unwrap();
12351        let header = CmfHeader {
12352            format: "cmf".into(),
12353            version: CMF_VERSION,
12354            arch,
12355            quant_type: QuantType::Q4Block,
12356            provenance: None,
12357            tokenizer_config: None,
12358            section_hashes: None,
12359            skills: Vec::new(),
12360            shard: None,
12361            calibration: None,
12362            routing: None,
12363        };
12364        let mut specs = Vec::new();
12365        for e in 0..ne {
12366            for (k, (n, r, c)) in [("g", inter, h), ("u", inter, h), ("d", h, inter)]
12367                .into_iter()
12368                .enumerate()
12369            {
12370                // Distinct experts: perturb only the nibble plane (any byte
12371                // is a valid pair of codes; the ladder stays intact).
12372                let mut data = synth_q4tp(r, c);
12373                for (i, byte) in data[..r * (c / GROUP_SIZE) * Q4TP_NIB]
12374                    .iter_mut()
12375                    .enumerate()
12376                {
12377                    *byte ^= ((i * (e * 3 + k + 1)) % 251) as u8;
12378                }
12379                specs.push(TensorSpec {
12380                    name: format!("{n}{e}"),
12381                    dtype: TensorDtype::Q4TiledP,
12382                    shape: vec![r, c],
12383                    data,
12384                });
12385            }
12386        }
12387        let dir = std::env::temp_dir().join(format!(
12388            "cmf-moe-rows-{}-{}",
12389            std::process::id(),
12390            FLOAT_ACTIVATIONS.get()
12391        ));
12392        std::fs::create_dir_all(&dir).unwrap();
12393        let path = dir.join("m.cmf");
12394        CmfModel::write(&path, &header, &specs, None, None).unwrap();
12395        let model = Arc::new(CmfModel::open(&path).unwrap());
12396        let t = |n: String| QTensor::from_model(&model, &n).unwrap();
12397        let g: Vec<QTensor> = (0..ne).map(|e| t(format!("g{e}"))).collect();
12398        let u: Vec<QTensor> = (0..ne).map(|e| t(format!("u{e}"))).collect();
12399        let d: Vec<QTensor> = (0..ne).map(|e| t(format!("d{e}"))).collect();
12400        let b = 4usize;
12401        let mut xs: Vec<f32> = (0..b * h)
12402            .map(|i| ((i * 31 + 7) % 89) as f32 / 89.0 - 0.5)
12403            .collect();
12404        xs[5] = 9.0; // an activation outlier on token 0
12405        // Token -> (experts in route order, weights).
12406        let routes: Vec<(Vec<usize>, Vec<f32>)> = vec![
12407            (vec![2, 0], vec![0.6, 0.4]),
12408            (vec![0, 1, 2], vec![0.2, 0.5, 0.3]),
12409            (vec![1], vec![1.0]),
12410            (vec![2, 1, 0], vec![0.25, 0.25, 0.5]),
12411        ];
12412        let pool = Pool::new(3);
12413        // Decode reference, token by token.
12414        let mut want = vec![0f32; b * h];
12415        for (tk, (idx, w)) in routes.iter().enumerate() {
12416            let x = &xs[tk * h..(tk + 1) * h];
12417            let pairs: Vec<(&QTensor, &QTensor)> = idx.iter().map(|&e| (&g[e], &u[e])).collect();
12418            let mut gs: Vec<Vec<f32>> = idx.iter().map(|_| vec![0f32; inter]).collect();
12419            assert!(QTensor::moe_gate_up_many(&pairs, x, &mut gs, Some(&pool)));
12420            if FLOAT_ACTIVATIONS.get() {
12421                for (slot, &e) in idx.iter().enumerate() {
12422                    let (mut gate, mut up) = (vec![0.0; inter], vec![0.0; inter]);
12423                    g[e].matvec(x, &mut gate, Some(&pool));
12424                    u[e].matvec(x, &mut up, Some(&pool));
12425                    for (v, u) in gate.iter_mut().zip(up) {
12426                        *v = (*v / (1.0 + (-*v).exp())) * u;
12427                    }
12428                    assert_eq!(gs[slot], gate, "float gate/up must equal ordinary matvecs");
12429                }
12430            }
12431            let downs: Vec<&QTensor> = idx.iter().map(|&e| &d[e]).collect();
12432            assert!(QTensor::moe_down_many(
12433                &downs,
12434                &gs,
12435                w,
12436                &mut want[tk * h..(tk + 1) * h],
12437                Some(&pool)
12438            ));
12439        }
12440        if FLOAT_ACTIVATIONS.get() {
12441            for (tk, (idx, w)) in routes.iter().enumerate() {
12442                let mut scalar = vec![0.0; h];
12443                for (&e, &weight) in idx.iter().zip(w) {
12444                    let (mut gate, mut up, mut down) =
12445                        (vec![0.0; inter], vec![0.0; inter], vec![0.0; h]);
12446                    g[e].matvec(&xs[tk * h..(tk + 1) * h], &mut gate, Some(&pool));
12447                    u[e].matvec(&xs[tk * h..(tk + 1) * h], &mut up, Some(&pool));
12448                    for (v, u) in gate.iter_mut().zip(up) {
12449                        *v = (*v / (1.0 + (-*v).exp())) * u;
12450                    }
12451                    d[e].matvec(&gate, &mut down, Some(&pool));
12452                    for (v, d) in scalar.iter_mut().zip(down) {
12453                        *v += weight * d;
12454                    }
12455                }
12456                assert_eq!(
12457                    &want[tk * h..(tk + 1) * h],
12458                    scalar,
12459                    "float many equals scalar experts"
12460                );
12461            }
12462        }
12463        // All four tokens at once, grouped by expert.
12464        let mut experts: Vec<usize> = Vec::new();
12465        let mut groups: Vec<Vec<usize>> = Vec::new();
12466        for (tk, (idx, _)) in routes.iter().enumerate() {
12467            for &e in idx {
12468                match experts.iter().position(|&x| x == e) {
12469                    Some(k) => groups[k].push(tk),
12470                    None => {
12471                        experts.push(e);
12472                        groups.push(vec![tk]);
12473                    }
12474                }
12475            }
12476        }
12477        let n_pairs: usize = groups.iter().map(|g| g.len()).sum();
12478        let pairs: Vec<(&QTensor, &QTensor)> = experts.iter().map(|&e| (&g[e], &u[e])).collect();
12479        let mut gs: Vec<Vec<f32>> = (0..n_pairs).map(|_| vec![0f32; inter]).collect();
12480        assert!(QTensor::moe_gate_up_rows(
12481            &pairs,
12482            &groups,
12483            &xs,
12484            &mut gs,
12485            Some(&pool)
12486        ));
12487        let downs: Vec<&QTensor> = experts.iter().map(|&e| &d[e]).collect();
12488        let lens: Vec<usize> = groups.iter().map(|g| g.len()).collect();
12489        let mut ds: Vec<Vec<f32>> = (0..n_pairs).map(|_| vec![0f32; h]).collect();
12490        assert!(QTensor::moe_down_rows(
12491            &downs,
12492            &lens,
12493            &gs,
12494            &mut ds,
12495            Some(&pool)
12496        ));
12497        let slot = |tk: usize, e: usize| {
12498            let k = experts.iter().position(|&x| x == e).unwrap();
12499            groups[..k].iter().map(|g| g.len()).sum::<usize>()
12500                + groups[k].iter().position(|&x| x == tk).unwrap()
12501        };
12502        let mut got = vec![0f32; b * h];
12503        for (tk, (idx, w)) in routes.iter().enumerate() {
12504            for i in 0..h {
12505                let mut acc = 0f32;
12506                for (&e, &we) in idx.iter().zip(w) {
12507                    acc += we * ds[slot(tk, e)][i];
12508                }
12509                got[tk * h + i] = acc;
12510            }
12511        }
12512        assert!(want.iter().any(|v| *v != 0.0));
12513        assert_eq!(
12514            want.iter().map(|v| v.to_bits()).collect::<Vec<_>>(),
12515            got.iter().map(|v| v.to_bits()).collect::<Vec<_>>(),
12516            "multi-token MoE must equal decode bit for bit"
12517        );
12518
12519        // Row-exact q4tp_matmat: five tokens (a blocked 1x4 tile + tail
12520        // otherwise) equal five matvecs.
12521        let b5 = 5usize;
12522        let x5: Vec<f32> = (0..b5 * h)
12523            .map(|i| ((i * 13 + 5) % 71) as f32 / 71.0 - 0.5)
12524            .collect();
12525        let mut mm = vec![0f32; b5 * inter];
12526        row_exact_scope(|| g[1].matmat(&x5, b5, &mut mm, Some(&pool)));
12527        for tk in 0..b5 {
12528            let mut mv = vec![0f32; inter];
12529            g[1].matvec(&x5[tk * h..(tk + 1) * h], &mut mv, Some(&pool));
12530            assert_eq!(
12531                mv.iter().map(|v| v.to_bits()).collect::<Vec<_>>(),
12532                mm[tk * inter..(tk + 1) * inter]
12533                    .iter()
12534                    .map(|v| v.to_bits())
12535                    .collect::<Vec<_>>(),
12536                "row-exact matmat token {tk}"
12537            );
12538        }
12539        // Other concurrent tests/requests may still hold the shared mode.
12540        // Nested, overlapping and unwind restoration is checked separately.
12541        let _ = std::fs::remove_dir_all(&dir);
12542    }
12543
12544    #[test]
12545    fn row_exact_scopes_survive_overlap_nesting_and_unwind() {
12546        use std::sync::{Barrier, atomic::{AtomicUsize, Ordering}};
12547        // A private counter makes this restoration test independent of
12548        // numerical tests concurrently using the production counter.
12549        let active = AtomicUsize::new(0);
12550        counted_row_exact_scope(&active, || {
12551            assert_eq!(active.load(Ordering::Acquire), 1);
12552            counted_row_exact_scope(&active, || {
12553                assert_eq!(active.load(Ordering::Acquire), 2);
12554            });
12555            assert_eq!(active.load(Ordering::Acquire), 1);
12556        });
12557        assert_eq!(active.load(Ordering::Acquire), 0);
12558
12559        let both_entered = Barrier::new(2);
12560        let release_last = Barrier::new(2);
12561        std::thread::scope(|s| {
12562            let first = s.spawn(|| counted_row_exact_scope(&active, || {
12563                both_entered.wait();
12564            }));
12565            let last = s.spawn(|| counted_row_exact_scope(&active, || {
12566                both_entered.wait();
12567                release_last.wait();
12568            }));
12569            first.join().unwrap();
12570            let after_first = active.load(Ordering::Acquire);
12571            release_last.wait();
12572            last.join().unwrap();
12573            assert_eq!(after_first, 1, "second request must remain exact");
12574        });
12575        assert_eq!(active.load(Ordering::Acquire), 0);
12576        let panic = std::panic::catch_unwind(|| {
12577            counted_row_exact_scope(&active, || panic!("scope unwind"));
12578        });
12579        assert!(panic.is_err());
12580        assert_eq!(active.load(Ordering::Acquire), 0);
12581    }
12582
12583    #[test]
12584    #[cfg(target_arch = "x86_64")]
12585    fn q4tp_float_avx2_is_bitwise_scalar() {
12586        if !avx2_enabled() {
12587            return;
12588        }
12589        for cols in [32, 64, 96, 2048, 4096] {
12590            let rows = 9;
12591            let bytes = synth_q4tp(rows, cols);
12592            let v = Q4tpView::new(&bytes, rows, cols);
12593            let gpr = cols / GROUP_SIZE;
12594            let mut sc = vec![0.0; gpr];
12595            for seed in 1..=5 {
12596                let xs: Vec<f32> = (0..cols)
12597                    .map(|i| (((i * 104729 + seed * 8191) % 100003) as f32 - 50001.0) / 7919.0)
12598                    .collect();
12599                for r in 0..rows {
12600                    v.scales_into(r, gpr, &mut sc);
12601                    let scalar = q4tp_row_float_scalar(v.nib, r, gpr, &xs, &sc);
12602                    let vector = unsafe { q4tp_row_float_avx2(v.nib, r, gpr, &xs, &sc) };
12603                    assert_eq!(
12604                        scalar.to_bits(),
12605                        vector.to_bits(),
12606                        "cols={cols} row={r} seed={seed}"
12607                    );
12608                }
12609            }
12610        }
12611    }
12612
12613    #[test]
12614    fn multi_token_moe_rows_float_equal_single_token_decode() {
12615        float_activations_scope(multi_token_moe_rows_equal_single_token_decode);
12616    }
12617
12618    #[test]
12619    fn full_gpu_q8_scope_is_nested_and_thread_local() {
12620        assert!(!FULL_GPU_Q8.get());
12621        let before = gpu_split_frac();
12622        {
12623            let _guard = enter_full_gpu_q8_scope();
12624            assert_eq!(gpu_split_frac(), 1.0);
12625            {
12626                let _nested = enter_full_gpu_q8_scope();
12627            }
12628            assert_eq!(gpu_split_frac(), 1.0);
12629            std::thread::spawn(|| assert!(!FULL_GPU_Q8.get())).join().unwrap();
12630        }
12631        assert!(!FULL_GPU_Q8.get());
12632        assert_eq!(gpu_split_frac(), before);
12633    }
12634
12635    #[test]
12636    fn float_activation_scope_is_nested_thread_local_and_unwind_safe() {
12637        assert!(!FLOAT_ACTIVATIONS.get());
12638        let before = a8w8_enabled();
12639        float_activations_scope(|| {
12640            assert!(!a8w8_enabled());
12641            float_activations_scope(|| assert!(!a8w8_enabled()));
12642            assert!(FLOAT_ACTIVATIONS.get());
12643            std::thread::spawn(|| assert!(!FLOAT_ACTIVATIONS.get()))
12644                .join()
12645                .unwrap();
12646        });
12647        assert!(!FLOAT_ACTIVATIONS.get());
12648        assert_eq!(a8w8_enabled(), before);
12649        let _ = std::panic::catch_unwind(|| float_activations_scope(|| panic!("test unwind")));
12650        assert!(!FLOAT_ACTIVATIONS.get());
12651    }
12652
12653    /// Batched q4/vbit matmat must equal per-position matvec calls
12654    /// exactly (the fallback it replaced) — same kernels, same order.
12655    #[test]
12656    fn batched_matmat_equals_per_position_matvec() {
12657        let (rows, cols, b) = (8, 64, 5);
12658        // q4 blob.
12659        let groups = rows * cols / GROUP_SIZE;
12660        let mut q4 = Vec::new();
12661        for i in 0..groups * 16 {
12662            q4.push((((i * 7 + 3) % 256) & 0xFF) as u8);
12663        }
12664        for g in 0..groups {
12665            q4.extend_from_slice(
12666                &cortiq_core::quant::f32_to_f16(0.01 + 0.003 * g as f32).to_le_bytes(),
12667            );
12668        }
12669        // vbit blob (mixed widths incl. 8).
12670        let ng = cols / GROUP_SIZE;
12671        let bits: Vec<u8> = vec![3, 4, 5, 6, 8, 4, 5, 3];
12672        let mut vb = bits.clone();
12673        for g in 0..rows * ng {
12674            vb.extend_from_slice(
12675                &cortiq_core::quant::f32_to_f16(0.02 + 0.001 * g as f32).to_le_bytes(),
12676            );
12677        }
12678        for r in 0..rows {
12679            let bw = bits[r] as usize;
12680            let (mut acc, mut nb) = (0u64, 0usize);
12681            let mut rowbytes = Vec::new();
12682            for i in 0..cols {
12683                let v = ((i * 7 + r * 13) % (1 << bw)) as u64;
12684                acc = (acc << bw) | v;
12685                nb += bw;
12686                while nb >= 8 {
12687                    nb -= 8;
12688                    rowbytes.push(((acc >> nb) & 0xFF) as u8);
12689                }
12690            }
12691            if nb > 0 {
12692                rowbytes.push(((acc << (8 - nb)) & 0xFF) as u8);
12693            }
12694            vb.extend_from_slice(&rowbytes);
12695        }
12696        let offsets = vbit_row_offsets(&vb, rows, cols);
12697
12698        let xs: Vec<f32> = (0..b * cols).map(|i| (i as f32 * 0.13).sin()).collect();
12699
12700        // q4: batch vs singles.
12701        let mut got = vec![0f32; b * rows];
12702        q4matmat(&q4, &xs, b, rows, cols, &mut got, None);
12703        for bi in 0..b {
12704            let mut expect = vec![0f32; rows];
12705            q4matvec(
12706                &q4,
12707                &xs[bi * cols..(bi + 1) * cols],
12708                rows,
12709                cols,
12710                &mut expect,
12711                None,
12712            );
12713            assert_eq!(
12714                &got[bi * rows..(bi + 1) * rows],
12715                &expect[..],
12716                "q4 batch pos {bi}"
12717            );
12718        }
12719
12720        // vbit: batch vs singles.
12721        let mut got = vec![0f32; b * rows];
12722        vbitmatmat(&vb, &offsets, &xs, b, rows, cols, &mut got, None);
12723        for bi in 0..b {
12724            let mut expect = vec![0f32; rows];
12725            vbitmatvec(
12726                &vb,
12727                &offsets,
12728                &xs[bi * cols..(bi + 1) * cols],
12729                rows,
12730                cols,
12731                &mut expect,
12732                None,
12733            );
12734            assert_eq!(
12735                &got[bi * rows..(bi + 1) * rows],
12736                &expect[..],
12737                "vbit batch pos {bi}"
12738            );
12739        }
12740    }
12741
12742    /// q4_tiled kernels must produce BIT-identical outputs to the q4
12743    /// split kernels on the same values (same ints, same order — only
12744    /// the byte placement differs).
12745    #[test]
12746    fn q4_tiled_matches_q4_block_bitexact() {
12747        let (rows, cols, b) = (8usize, 128usize, 3usize);
12748        let groups = rows * cols / GROUP_SIZE;
12749        let mut split = Vec::with_capacity(groups * 18);
12750        for i in 0..groups * 16 {
12751            split.push((((i * 7 + 3) % 256) & 0xFF) as u8);
12752        }
12753        for g in 0..groups {
12754            split.extend_from_slice(
12755                &cortiq_core::quant::f32_to_f16(0.01 + 0.003 * g as f32).to_le_bytes(),
12756            );
12757        }
12758        // Re-tile: [scale][nibbles] per group.
12759        let (packed, scales) = split.split_at(groups * 16);
12760        let mut tiled = Vec::with_capacity(groups * Q4_TILE);
12761        for g in 0..groups {
12762            tiled.extend_from_slice(&scales[g * 2..g * 2 + 2]);
12763            tiled.extend_from_slice(&packed[g * 16..(g + 1) * 16]);
12764        }
12765
12766        let mut x1: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.17).sin()).collect();
12767        x1[9] = 250.0; // exercise the outlier path
12768        let x2: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.23).cos()).collect();
12769
12770        let (mut a, mut t) = (vec![0f32; rows], vec![0f32; rows]);
12771        q4matvec(&split, &x1, rows, cols, &mut a, None);
12772        q4t_matvec(&tiled, &x1, rows, cols, &mut t, None);
12773        assert_eq!(a, t, "q4t matvec must match q4 bit-for-bit");
12774
12775        let (mut a1, mut a2) = (vec![0f32; rows], vec![0f32; rows]);
12776        let (mut t1, mut t2) = (vec![0f32; rows], vec![0f32; rows]);
12777        q4matvec2(&split, &x1, &x2, rows, cols, &mut a1, &mut a2, None);
12778        q4t_matvec2(&tiled, &x1, &x2, rows, cols, &mut t1, &mut t2, None);
12779        assert_eq!(a1, t1);
12780        assert_eq!(a2, t2);
12781
12782        let xs: Vec<f32> = (0..b * cols).map(|i| (i as f32 * 0.13).sin()).collect();
12783        let (mut am, mut tm) = (vec![0f32; b * rows], vec![0f32; b * rows]);
12784        q4matmat(&split, &xs, b, rows, cols, &mut am, None);
12785        q4t_matmat(&tiled, &xs, b, rows, cols, &mut tm, None);
12786        assert_eq!(am, tm, "q4t matmat must match q4 bit-for-bit");
12787    }
12788
12789    /// q4 SDOT outlier correction: a single huge activation channel
12790    /// (>8·rms → outlier, zeroed in xq) must still contribute its EXACT
12791    /// term. On-grid bulk (±1/0 → xq dequantizes exactly) isolates the
12792    /// correction from A8W8 noise. cols must exceed 64: at n=64 the
12793    /// 8·rms threshold equals sqrt(v²+rest) ≥ v, so a single outlier
12794    /// can never qualify (8² = n).
12795    #[test]
12796    fn q4matvec_sdot_outlier_exact() {
12797        let (rows, cols) = (4, 128);
12798        let groups = rows * cols / GROUP_SIZE;
12799        let mut bytes = Vec::with_capacity(groups * 16 + groups * 2);
12800        for i in 0..groups * 16 {
12801            bytes.push(((i * 11 + 5) % 256) as u8);
12802        }
12803        for g in 0..groups {
12804            let s = 0.02 + 0.002 * g as f32;
12805            bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
12806        }
12807        let mut x: Vec<f32> = (0..cols)
12808            .map(|i| match i % 3 {
12809                0 => 1.0,
12810                1 => -1.0,
12811                _ => 0.0,
12812            })
12813            .collect();
12814        x[17] = 300.0; // ≫ 8·rms → outlier channel
12815
12816        let mut reference = vec![0.0f32; rows * cols];
12817        cortiq_core::quant::dequant_q4_block(&bytes, &mut reference);
12818        let mut expect = vec![0.0f32; rows];
12819        for r in 0..rows {
12820            expect[r] = reference[r * cols..(r + 1) * cols]
12821                .iter()
12822                .zip(&x)
12823                .map(|(w, xv)| w * xv)
12824                .sum();
12825        }
12826        let mut got = vec![0.0f32; rows];
12827        q4matvec(&bytes, &x, rows, cols, &mut got, None);
12828        let scale = expect.iter().fold(0f32, |m, v| m.max(v.abs())).max(1.0);
12829        for r in 0..rows {
12830            assert!(
12831                (got[r] - expect[r]).abs() < 2e-3 * scale,
12832                "row {r}: {} vs {} (outlier term must be exact)",
12833                got[r],
12834                expect[r]
12835            );
12836        }
12837    }
12838
12839    /// The fused q1t matvec must equal the reference (dequant_q1t → dot),
12840    /// including the ternary zero level and the binary-searched outlier
12841    /// overlay. Guards the mmap kernel that makes a 12B q1t runnable.
12842    #[test]
12843    fn q1t_matvec_matches_reference() {
12844        use cortiq_core::quant::{dequant_q1t, f32_to_f16};
12845        let (rows, cols) = (3usize, 64usize); // gpr = 2
12846        let gpr = cols / GROUP_SIZE;
12847        let scales = [0.5f32, 0.3, 0.7, 0.2, 0.6, 0.15];
12848        // Overlay (must be sorted by flat index): a few spikes across rows.
12849        let outliers: [(u32, f32); 3] = [(5, 9.0), (70, -4.5), (150, 3.25)];
12850        let is_out = |flat: usize| outliers.iter().any(|&(i, _)| i as usize == flat);
12851        let mut bytes = Vec::new();
12852        for r in 0..rows {
12853            for g in 0..gpr {
12854                bytes.extend_from_slice(&f32_to_f16(scales[r * gpr + g]).to_le_bytes());
12855                let mut c = [0u8; 7];
12856                for k in 0..GROUP_SIZE {
12857                    // Encoder invariant: code 0 at outlier positions.
12858                    let code = if is_out(r * cols + g * GROUP_SIZE + k) {
12859                        0
12860                    } else {
12861                        ((k + r * 3 + g) % 3) as u8 // 0,1,2
12862                    };
12863                    cortiq_core::quant::q1t_pack(&mut c, k, code);
12864                }
12865                bytes.extend_from_slice(&c);
12866            }
12867        }
12868        // Per-row overlay: [u32 row_ptr[rows+1]] then [(u16 col, f16 val)] by
12869        // row (outliers are sorted by flat index → already grouped by row).
12870        let mut row_ptr = vec![0u32; rows + 1];
12871        for &(idx, _) in &outliers {
12872            row_ptr[idx as usize / cols + 1] += 1;
12873        }
12874        for r in 0..rows {
12875            row_ptr[r + 1] += row_ptr[r];
12876        }
12877        for &p in &row_ptr {
12878            bytes.extend_from_slice(&p.to_le_bytes());
12879        }
12880        for &(idx, v) in &outliers {
12881            bytes.extend_from_slice(&((idx as usize % cols) as u16).to_le_bytes());
12882            bytes.extend_from_slice(&f32_to_f16(v).to_le_bytes());
12883        }
12884
12885        let mut refw = vec![0f32; rows * cols];
12886        dequant_q1t(&bytes, rows, cols, &mut refw);
12887        // On-grid activations (±1, amax 1) so the int8 SDOT path reconstructs
12888        // x exactly and matches the f32 reference (same trick as the q1 test).
12889        let x: Vec<f32> = (0..cols)
12890            .map(|j| if j % 3 == 0 { 1.0 } else { -1.0 })
12891            .collect();
12892        let mut expect = vec![0f32; rows];
12893        for r in 0..rows {
12894            let mut a = 0.0f32;
12895            for j in 0..cols {
12896                a += refw[r * cols + j] * x[j];
12897            }
12898            expect[r] = a;
12899        }
12900        let tol = |e: f32| 1e-3 * e.abs().max(1e-3);
12901        let mut got = vec![0f32; rows];
12902        q1t_matvec(&bytes, &x, rows, cols, &mut got, None);
12903        for r in 0..rows {
12904            assert!(
12905                (got[r] - expect[r]).abs() < tol(expect[r]),
12906                "row {r}: {} vs {}",
12907                got[r],
12908                expect[r]
12909            );
12910        }
12911        // matmat (b=2, f32 decode path) must agree too.
12912        let x2: Vec<f32> = x.iter().chain(x.iter()).copied().collect();
12913        let mut gm = vec![0f32; 2 * rows];
12914        q1t_matmat(&bytes, &x2, 2, rows, cols, &mut gm, None);
12915        for r in 0..rows {
12916            assert!((gm[r] - expect[r]).abs() < tol(expect[r]));
12917            assert!((gm[rows + r] - expect[r]).abs() < tol(expect[r]));
12918        }
12919        // Fused pair (q1t_matvec2) must equal two single matvecs
12920        // bit-for-bit: same unpack, same group order, same f32
12921        // accumulation per stream. Distinct x2 exercises both lanes.
12922        let xb: Vec<f32> = (0..cols)
12923            .map(|j| if j % 5 == 0 { -1.0 } else { 1.0 })
12924            .collect();
12925        let (mut s1, mut s2) = (vec![0f32; rows], vec![0f32; rows]);
12926        q1t_matvec(&bytes, &x, rows, cols, &mut s1, None);
12927        q1t_matvec(&bytes, &xb, rows, cols, &mut s2, None);
12928        let (mut p1, mut p2) = (vec![0f32; rows], vec![0f32; rows]);
12929        q1t_matvec2(&bytes, &x, &xb, rows, cols, &mut p1, &mut p2, None);
12930        assert_eq!(p1, s1, "q1t pair lane 1 ≠ single matvec");
12931        assert_eq!(p2, s2, "q1t pair lane 2 ≠ single matvec");
12932    }
12933
12934    /// Pair == 2×matvec with an ODD group count (the kernel's tail
12935    /// group) and no overlay section.
12936    #[test]
12937    fn q1t_matvec2_odd_gpr_matches_singles() {
12938        use cortiq_core::quant::{Q1T_TILE, f32_to_f16, q1t_pack};
12939        let (rows, cols) = (5usize, 96usize); // gpr = 3 → paired + tail
12940        let gpr = cols / GROUP_SIZE;
12941        let mut bytes = Vec::with_capacity(rows * gpr * Q1T_TILE);
12942        for r in 0..rows {
12943            for g in 0..gpr {
12944                bytes.extend_from_slice(&f32_to_f16(0.1 + 0.05 * (r + g) as f32).to_le_bytes());
12945                let mut c = [0u8; 7];
12946                for k in 0..GROUP_SIZE {
12947                    q1t_pack(&mut c, k, ((k * 7 + r * 5 + g * 3) % 3) as u8);
12948                }
12949                bytes.extend_from_slice(&c);
12950            }
12951        }
12952        let x1: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.31).sin()).collect();
12953        let x2: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.17).cos()).collect();
12954        let (mut s1, mut s2) = (vec![0f32; rows], vec![0f32; rows]);
12955        q1t_matvec(&bytes, &x1, rows, cols, &mut s1, None);
12956        q1t_matvec(&bytes, &x2, rows, cols, &mut s2, None);
12957        let (mut p1, mut p2) = (vec![0f32; rows], vec![0f32; rows]);
12958        q1t_matvec2(&bytes, &x1, &x2, rows, cols, &mut p1, &mut p2, None);
12959        assert_eq!(p1, s1, "odd-gpr pair lane 1 ≠ single");
12960        assert_eq!(p2, s2, "odd-gpr pair lane 2 ≠ single");
12961    }
12962
12963    // Speed A/B: fused pair (one unpack, two streams) vs two single
12964    // matvecs. Single-threaded, FFN-sized, min-of paired in-process.
12965    //   cargo test -p cortiq-engine --release q1t_matvec2_speed -- --ignored --nocapture
12966    #[test]
12967    #[ignore]
12968    fn q1t_matvec2_speed() {
12969        use cortiq_core::quant::{Q1T_TILE, f32_to_f16, q1t_pack};
12970        use std::time::Instant;
12971        let (rows, cols) = (8192usize, 4096usize);
12972        let gpr = cols / GROUP_SIZE;
12973        let mut bytes = Vec::with_capacity(rows * gpr * Q1T_TILE);
12974        for r in 0..rows {
12975            for g in 0..gpr {
12976                let s = 0.1 + ((r + g) % 7) as f32 * 0.01;
12977                bytes.extend_from_slice(&f32_to_f16(s).to_le_bytes());
12978                let mut c = [0u8; 7];
12979                for k in 0..GROUP_SIZE {
12980                    q1t_pack(&mut c, k, ((k * 7 + r + g) % 3) as u8);
12981                }
12982                bytes.extend_from_slice(&c);
12983            }
12984        }
12985        let x1: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.31).sin()).collect();
12986        let x2: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.17).cos()).collect();
12987        let (mut s1, mut s2) = (vec![0f32; rows], vec![0f32; rows]);
12988        let (mut p1, mut p2) = (vec![0f32; rows], vec![0f32; rows]);
12989        // Warm both paths once.
12990        q1t_matvec(&bytes, &x1, rows, cols, &mut s1, None);
12991        q1t_matvec2(&bytes, &x1, &x2, rows, cols, &mut p1, &mut p2, None);
12992        let (mut t_pair, mut t_two) = (f64::MAX, f64::MAX);
12993        for _ in 0..8 {
12994            let t0 = Instant::now();
12995            q1t_matvec2(&bytes, &x1, &x2, rows, cols, &mut p1, &mut p2, None);
12996            t_pair = t_pair.min(t0.elapsed().as_secs_f64() * 1000.0);
12997            let t1 = Instant::now();
12998            q1t_matvec(&bytes, &x1, rows, cols, &mut s1, None);
12999            q1t_matvec(&bytes, &x2, rows, cols, &mut s2, None);
13000            t_two = t_two.min(t1.elapsed().as_secs_f64() * 1000.0);
13001        }
13002        assert_eq!(p1, s1);
13003        assert_eq!(p2, s2);
13004        println!("q1t pair {rows}x{cols}: fused {t_pair:.2} ms | two singles {t_two:.2} ms");
13005    }
13006
13007    // Speed A/B: the base-3-division decode (what the packing commit left in
13008    // place) vs the fused sign-LUT matvec. Both single-threaded, same bytes.
13009    //   cargo test -p cortiq-engine q1t_matvec_speed -- --ignored --nocapture
13010    #[test]
13011    #[ignore]
13012    fn q1t_matvec_speed() {
13013        use cortiq_core::quant::{Q1T_TILE, f32_to_f16, q1t_code, q1t_pack};
13014        use std::time::Instant;
13015        let (rows, cols) = (8192usize, 4096usize); // FFN-sized
13016        let gpr = cols / GROUP_SIZE;
13017        let mut bytes = Vec::with_capacity(rows * gpr * Q1T_TILE + 16);
13018        for r in 0..rows {
13019            for g in 0..gpr {
13020                let s = 0.1 + ((r + g) % 7) as f32 * 0.01;
13021                bytes.extend_from_slice(&f32_to_f16(s).to_le_bytes());
13022                let mut c = [0u8; 7];
13023                for k in 0..GROUP_SIZE {
13024                    q1t_pack(&mut c, k, ((k * 7 + r + g) % 3) as u8);
13025                }
13026                bytes.extend_from_slice(&c);
13027            }
13028        }
13029        let (n, stride) = (rows * cols, 40usize); // ~2.5% outliers, per-row overlay
13030        let mut row_ptr = vec![0u32; rows + 1];
13031        let mut idx = 0usize;
13032        while idx < n {
13033            row_ptr[idx / cols + 1] += 1;
13034            idx += stride;
13035        }
13036        for r in 0..rows {
13037            row_ptr[r + 1] += row_ptr[r];
13038        }
13039        for &p in &row_ptr {
13040            bytes.extend_from_slice(&p.to_le_bytes());
13041        }
13042        let mut idx = 0usize;
13043        while idx < n {
13044            bytes.extend_from_slice(&((idx % cols) as u16).to_le_bytes());
13045            bytes.extend_from_slice(&f32_to_f16((idx % 13) as f32 * 0.1 - 0.6).to_le_bytes());
13046            idx += stride;
13047        }
13048        // On-grid ±1 so the fast path's int8 SDOT is exact vs the f32 "slow"
13049        // reference (the A/B is a timing check; values must still agree).
13050        let x: Vec<f32> = (0..cols)
13051            .map(|j| if j % 3 == 0 { 1.0 } else { -1.0 })
13052            .collect();
13053        let (rp_off, ent_off, has_ov) = q1t_overlay(&bytes, rows * gpr * Q1T_TILE, rows);
13054
13055        // "before": base-3 division decode into a buffer, then dot.
13056        let slow = |out: &mut [f32]| {
13057            let mut buf = vec![0f32; cols];
13058            for r in 0..rows {
13059                for g in 0..gpr {
13060                    let off = (r * gpr + g) * Q1T_TILE;
13061                    let s = f16_to_f32(u16::from_le_bytes([bytes[off], bytes[off + 1]]));
13062                    let codes = &bytes[off + 2..off + Q1T_TILE];
13063                    for k in 0..GROUP_SIZE {
13064                        buf[g * GROUP_SIZE + k] = match q1t_code(codes, k) {
13065                            1 => s,
13066                            2 => -s,
13067                            _ => 0.0,
13068                        };
13069                    }
13070                }
13071                out[r] = q1t_row_outlier_correction(&bytes, r, rp_off, ent_off, has_ov, &x)
13072                    + (0..cols).map(|j| buf[j] * x[j]).sum::<f32>();
13073            }
13074        };
13075        let iters = 5;
13076        let mut a = vec![0f32; rows];
13077        slow(&mut a); // warm
13078        let t = Instant::now();
13079        for _ in 0..iters {
13080            slow(&mut a);
13081        }
13082        let slow_ms = t.elapsed().as_secs_f64() * 1e3 / iters as f64;
13083
13084        let mut b = vec![0f32; rows];
13085        q1t_matvec(&bytes, &x, rows, cols, &mut b, None); // warm
13086        let t = Instant::now();
13087        for _ in 0..iters {
13088            q1t_matvec(&bytes, &x, rows, cols, &mut b, None);
13089        }
13090        let fast_ms = t.elapsed().as_secs_f64() * 1e3 / iters as f64;
13091
13092        for r in 0..rows {
13093            assert!((a[r] - b[r]).abs() < 1e-2, "mismatch row {r}");
13094        }
13095        println!(
13096            "q1t matvec {rows}x{cols} (1 thread): div-decode {slow_ms:.2} ms  fused-LUT {fast_ms:.2} ms  => {:.2}x",
13097            slow_ms / fast_ms
13098        );
13099    }
13100}
13101
13102#[cfg(test)]
13103mod gemm_bench {
13104    /// `cargo test -p cortiq-engine --release q4tp_matmat_throughput -- --ignored --nocapture`
13105    /// Times the batched q4tp GEMM at the shapes the image DiT runs
13106    /// (b=296 tokens, 2304 -> 9216), on synthetic bytes: no model, no
13107    /// mmap, no thermal drift over minutes — a kernel change shows up
13108    /// here in seconds where a full render hides it in noise.
13109    ///
13110    /// On macOS add `CMF_ACCEL=0`: this shape is over the 500k-cell mark
13111    /// where the matmat hands off to Accelerate's dequant sgemm, and
13112    /// without the opt-out both rows below measure the AMX, not the
13113    /// kernel under test.
13114    #[test]
13115    #[ignore]
13116    fn q4tp_matmat_throughput() {
13117        // 296 is a prompt-encode batch; the image DiT runs 2085 at
13118        // 512x512, where the activation panel stops fitting L2 and the
13119        // loop's shape starts to matter more than its instructions.
13120        let b: usize = std::env::var("CMF_BENCH_B")
13121            .ok()
13122            .and_then(|v| v.parse().ok())
13123            .unwrap_or(296);
13124        let (rows, cols) = (9216usize, 2304usize);
13125        let (_, _, _) = (rows, cols, b);
13126        let total =
13127            cortiq_core::quant::expected_nbytes(cortiq_core::TensorDtype::Q4TiledP, &[rows, cols])
13128                .unwrap();
13129        // Random nibbles are fine, but the row params are f16 (lo, step)
13130        // of a geometric ladder: garbage there gives exp2 of a huge
13131        // exponent, the scales come back inf, and the whole bench times
13132        // NaN arithmetic instead of the kernel.
13133        let (params_off, codes_off, _) = cortiq_core::quant::q4tp_sections(rows, cols);
13134        let mut bytes: Vec<u8> = (0..total).map(|i| (i * 37 % 251) as u8).collect();
13135        let lo = cortiq_core::quant::f32_to_f16(-4.0);
13136        let step = cortiq_core::quant::f32_to_f16(0.1);
13137        for r in 0..rows {
13138            let o = params_off + r * 4;
13139            bytes[o..o + 2].copy_from_slice(&lo.to_le_bytes());
13140            bytes[o + 2..o + 4].copy_from_slice(&step.to_le_bytes());
13141        }
13142        let _ = codes_off;
13143        let xs: Vec<f32> = (0..b * cols)
13144            .map(|i| ((i % 97) as f32 - 48.0) / 48.0)
13145            .collect();
13146        let mut out = vec![0f32; b * rows];
13147        let pool = crate::pool::Pool::from_env();
13148        // A shared 48-core stand drifts ±25% run to run, which is wider
13149        // than any kernel change worth making. So: alternate the two
13150        // kernels inside one process and keep the BEST time for
13151        // each. Interleaving makes both see the same interference, and a
13152        // minimum is the one statistic another tenant cannot inflate.
13153        super::q4tp_matmat(&bytes, &xs, b, rows, cols, &mut out, pool.as_deref());
13154        let reps: usize = std::env::var("CMF_BENCH_REPS")
13155            .ok()
13156            .and_then(|v| v.parse().ok())
13157            .unwrap_or(10);
13158        let mut best = [f64::MAX; 2];
13159        let mut sums = [0f32; 2];
13160        for _ in 0..reps {
13161            for (k, w) in [(0usize, 1u8), (1usize, 2u8)] {
13162                super::Q4TP_ALT.store(w, std::sync::atomic::Ordering::Relaxed);
13163                let t = std::time::Instant::now();
13164                super::q4tp_matmat(&bytes, &xs, b, rows, cols, &mut out, pool.as_deref());
13165                best[k] = best[k].min(t.elapsed().as_secs_f64());
13166                sums[k] = out.iter().take(64).sum::<f32>();
13167            }
13168        }
13169        let flops = 2.0 * b as f64 * rows as f64 * cols as f64;
13170        for (k, name) in ["previous", "tuned   "].iter().enumerate() {
13171            println!(
13172                "q4tp matmat {rows}x{cols} b={b} {name}: {:.1} ms  {:.1} GFLOP/s  (checksum {:.3})",
13173                best[k] * 1e3,
13174                flops / best[k] / 1e9,
13175                sums[k]
13176            );
13177        }
13178        assert!(
13179            (sums[0] - sums[1]).abs() < 1e-2,
13180            "the tuned kernel changed the result: {} vs {}",
13181            sums[0],
13182            sums[1]
13183        );
13184    }
13185
13186    /// The blocked kernel must agree with the per-column path exactly —
13187    /// same weights, same activation split, only a different instruction
13188    /// mix. Shapes are chosen to hit the awkward cases: a column count
13189    /// that leaves an odd group (the 512-bit kernel does two at a time),
13190    /// and a batch that does not divide by four.
13191    #[test]
13192    fn q4tp_matmat_blocked_matches_scalar() {
13193        use std::sync::atomic::Ordering::Relaxed;
13194        // The last shape carries the image DiT's column count — 2304, so
13195        // 72 groups of accumulation, which is where a reordered sum can
13196        // actually drift — and runs through the thread pool, since the
13197        // blocked path splits rows across workers. Its row count stays
13198        // under 500k cells on purpose: above that, macOS diverts the whole
13199        // matmat to the Accelerate/AMX dequant sgemm and neither kernel
13200        // here would run.
13201        for &(rows, cols, b) in &[
13202            (64usize, 128usize, 7usize),
13203            (33, 96, 4),
13204            (16, 256, 9),
13205            (192, 2304, 37),
13206        ] {
13207            let total = cortiq_core::quant::expected_nbytes(
13208                cortiq_core::TensorDtype::Q4TiledP,
13209                &[rows, cols],
13210            )
13211            .unwrap();
13212            let (params_off, _, _) = cortiq_core::quant::q4tp_sections(rows, cols);
13213            let mut bytes: Vec<u8> = (0..total).map(|i| (i * 61 % 251) as u8).collect();
13214            let lo = cortiq_core::quant::f32_to_f16(-4.0);
13215            let step = cortiq_core::quant::f32_to_f16(0.1);
13216            for r in 0..rows {
13217                let o = params_off + r * 4;
13218                bytes[o..o + 2].copy_from_slice(&lo.to_le_bytes());
13219                bytes[o + 2..o + 4].copy_from_slice(&step.to_le_bytes());
13220            }
13221            let xs: Vec<f32> = (0..b * cols)
13222                .map(|i| ((i % 89) as f32 - 44.0) / 44.0)
13223                .collect();
13224            let mut got = vec![0f32; b * rows];
13225            let mut want = vec![0f32; b * rows];
13226            let gpr = cols / 32;
13227            let view = super::Q4tpView::new(&bytes, rows, cols);
13228            let pool = crate::pool::Pool::from_env();
13229            super::Q4TP_ALT.store(2, Relaxed);
13230            super::q4tp_matmat(&bytes, &xs, b, rows, cols, &mut got, pool.as_deref());
13231            super::Q4TP_ALT.store(1, Relaxed);
13232            super::q4tp_matmat(&bytes, &xs, b, rows, cols, &mut want, pool.as_deref());
13233            super::Q4TP_ALT.store(0, Relaxed);
13234            // Measured against the output's scale, not cell by cell: a
13235            // dot product of 2304 terms lands near zero wherever the row
13236            // and the activation nearly cancel, and there a per-cell
13237            // ratio reports 1e-3 for an absolute error of 5e-6 — f32's
13238            // own rounding, reordered. What must stay small is the error
13239            // relative to what the layer actually outputs.
13240            let scale = want.iter().fold(0f32, |m, v| m.max(v.abs())).max(1e-6);
13241            let (mut worst, mut at) = (0f32, 0usize);
13242            for (i, (g, w)) in got.iter().zip(&want).enumerate() {
13243                if (g - w).abs() > worst {
13244                    worst = (g - w).abs();
13245                    at = i;
13246                }
13247            }
13248            assert!(
13249                worst <= 1e-4 * scale,
13250                "{rows}x{cols} b={b}: blocked and scalar disagree by {worst:.3e} \
13251                 (scale {scale:.3e}) at cell {at}: {} vs {}",
13252                got[at],
13253                want[at]
13254            );
13255
13256            // "Same speed, no quality loss" is a claim about which answer
13257            // is RIGHT, not about which two agree. Both paths sum the same
13258            // 2304 products in different orders, so f64 decides: the
13259            // blocked kernel keeps sixteen partial sums and folds them at
13260            // the end, which is a shallower addition tree than the
13261            // per-column path's running scalar, and it must not be worse.
13262            let (mut e_blocked, mut e_scalar) = (0f64, 0f64);
13263            for bi in 0..b {
13264                let act = super::split_act(&xs[bi * cols..(bi + 1) * cols]);
13265                for r in 0..rows {
13266                    let mut sc = vec![0f32; gpr];
13267                    view.scales_into(r, gpr, &mut sc);
13268                    let mut exact = 0f64;
13269                    for j in 0..cols {
13270                        let (w, sq) = super::q4tp_outlier(view.nib, r, gpr, j, &sc);
13271                        exact += w as f64 * sq as f64 * act.xq[j] as f64;
13272                    }
13273                    exact *= act.sx as f64;
13274                    for &(j, xv) in &act.outliers {
13275                        let (w, sq) = super::q4tp_outlier(view.nib, r, gpr, j, &sc);
13276                        exact += w as f64 * sq as f64 * xv as f64;
13277                    }
13278                    let i = bi * rows + r;
13279                    e_blocked = e_blocked.max((got[i] as f64 - exact).abs());
13280                    e_scalar = e_scalar.max((want[i] as f64 - exact).abs());
13281                }
13282            }
13283            println!(
13284                "{rows}x{cols} b={b}: worst error vs f64 — blocked {e_blocked:.3e}, \
13285                 per-column {e_scalar:.3e}"
13286            );
13287            // An absolute bar, not a race between the two: at these
13288            // magnitudes both sit in f32's last bits, and on a small shape
13289            // whichever one happens to round the unluckiest cell "wins" by
13290            // a factor the next seed reverses.
13291            assert!(
13292                e_blocked <= 1e-5 * scale as f64 && e_scalar <= 1e-5 * scale as f64,
13293                "{rows}x{cols} b={b}: error against f64 too large — blocked \
13294                 {e_blocked:.3e}, per-column {e_scalar:.3e}, scale {scale:.3e}"
13295            );
13296        }
13297    }
13298
13299    /// The q4t twin of the throughput bench, same shape and rules, so the
13300    /// two quantisations' batch kernels can be read against each other.
13301    /// `cargo test -p cortiq-engine --release q4t_matmat_throughput -- --ignored --nocapture`
13302    #[test]
13303    #[ignore]
13304    fn q4t_matmat_throughput() {
13305        let (rows, cols, b) = (9216usize, 2304usize, 296usize);
13306        let total =
13307            cortiq_core::quant::expected_nbytes(cortiq_core::TensorDtype::Q4Tiled, &[rows, cols])
13308                .unwrap();
13309        // q4t carries a per-group f16 scale in the tile's first two bytes;
13310        // random bytes there decode to inf and the bench would time NaNs.
13311        let mut bytes: Vec<u8> = (0..total).map(|i| (i * 37 % 251) as u8).collect();
13312        let sc = cortiq_core::quant::f32_to_f16(0.02);
13313        for t in bytes.chunks_mut(super::Q4_TILE) {
13314            t[..2].copy_from_slice(&sc.to_le_bytes());
13315        }
13316        let xs: Vec<f32> = (0..b * cols)
13317            .map(|i| ((i % 97) as f32 - 48.0) / 48.0)
13318            .collect();
13319        let mut out = vec![0f32; b * rows];
13320        let pool = crate::pool::Pool::from_env();
13321        super::q4t_matmat(&bytes, &xs, b, rows, cols, &mut out, pool.as_deref());
13322        let reps: usize = std::env::var("CMF_BENCH_REPS")
13323            .ok()
13324            .and_then(|v| v.parse().ok())
13325            .unwrap_or(10);
13326        let mut best = f64::MAX;
13327        for _ in 0..reps {
13328            let t = std::time::Instant::now();
13329            super::q4t_matmat(&bytes, &xs, b, rows, cols, &mut out, pool.as_deref());
13330            best = best.min(t.elapsed().as_secs_f64());
13331        }
13332        let flops = 2.0 * b as f64 * rows as f64 * cols as f64;
13333        println!(
13334            "q4t matmat {rows}x{cols} b={b}: {:.1} ms  {:.1} GFLOP/s  (checksum {:.3})",
13335            best * 1e3,
13336            flops / best / 1e9,
13337            out.iter().take(64).sum::<f32>()
13338        );
13339    }
13340}