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
116impl QTensor {
117    pub fn from_f32(data: Vec<f32>, rows: usize, cols: usize) -> Self {
118        debug_assert_eq!(data.len(), rows * cols);
119        Self::F32 { data, rows, cols }
120    }
121
122    /// Wrap a directory tensor without dequantizing the payload.
123    /// Falls back to dequantized f32 for dtypes without a fused kernel.
124    pub fn from_model(model: &Arc<CmfModel>, name: &str) -> Result<Self, String> {
125        // Indexed lookup: the linear directory scan made pipeline build
126        // O(N²) on MoE/skills files with thousands of tensors.
127        let idx = model
128            .tensor_index(name)
129            .ok_or_else(|| format!("tensor '{name}' not found in CMF directory"))?;
130        let entry = &model.tensors[idx];
131        if entry.shape.len() != 2 {
132            return Err(format!("QTensor::from_model needs 2-D, got '{name}'"));
133        }
134        let (rows, cols) = (entry.shape[0], entry.shape[1]);
135        let bytes = model.entry_bytes(entry);
136
137        match entry.dtype {
138            TensorDtype::Q8Row | TensorDtype::Q8_2f => {
139                let n = rows * cols;
140                let scales_off = n;
141                let row_scale: Vec<f32> = (0..rows)
142                    .map(|o| {
143                        f16_to_f32(u16::from_le_bytes([
144                            bytes[scales_off + o * 2],
145                            bytes[scales_off + o * 2 + 1],
146                        ]))
147                    })
148                    .collect();
149                let col_field: Vec<f32> = if entry.dtype == TensorDtype::Q8_2f {
150                    let col_off = n + rows * 2;
151                    (0..cols)
152                        .map(|i| {
153                            f16_to_f32(u16::from_le_bytes([
154                                bytes[col_off + i * 2],
155                                bytes[col_off + i * 2 + 1],
156                            ]))
157                        })
158                        .collect()
159                } else {
160                    Vec::new()
161                };
162                Ok(Self::Mapped {
163                    model: model.clone(),
164                    idx,
165                    dtype: entry.dtype,
166                    rows,
167                    cols,
168                    row_scale,
169                    col_field,
170                    vbit_offsets: Vec::new(),
171                    repack: q8_repack(bytes, rows, cols),
172                })
173            }
174            // vbit: fused kernel unpacks variable-bit rows from mmap.
175            TensorDtype::Vbit if cols % GROUP_SIZE == 0 => Ok(Self::Mapped {
176                model: model.clone(),
177                idx,
178                dtype: entry.dtype,
179                rows,
180                cols,
181                row_scale: Vec::new(),
182                col_field: Vec::new(),
183                vbit_offsets: vbit_row_offsets(bytes, rows, cols),
184                repack: Vec::new(),
185            }),
186            // vbit_ro (§4.2): the offset table comes straight from the
187            // file — no load-time prefix scan; kernels are shared with
188            // legacy vbit (they consume absolute offsets either way).
189            TensorDtype::VbitRo if cols % GROUP_SIZE == 0 => {
190                let (_, off_off, packed_off) = cortiq_core::quant::vbit_ro_sections(rows, cols);
191                let offsets: Vec<usize> = (0..=rows)
192                    .map(|r| packed_off + cortiq_core::quant::vbit_ro_offset(bytes, off_off, r))
193                    .collect();
194                Ok(Self::Mapped {
195                    model: model.clone(),
196                    idx,
197                    dtype: entry.dtype,
198                    rows,
199                    cols,
200                    row_scale: Vec::new(),
201                    col_field: Vec::new(),
202                    vbit_offsets: offsets,
203                    repack: Vec::new(),
204                })
205            }
206            // q4_block: fused kernel reads nibbles straight from mmap —
207            // a 14B q4 file no longer explodes into ×8 f32 RAM.
208            // q4_tiled (§4.3): interleaved [scale][nibbles] tiles — one
209            // sequential memory stream (measured ×1.66 ARM / ×1.13 AVX2
210            // at kernel level over the split layout).
211            TensorDtype::Q4Tiled if cols % GROUP_SIZE == 0 => Ok(Self::Mapped {
212                model: model.clone(),
213                idx,
214                dtype: entry.dtype,
215                rows,
216                cols,
217                row_scale: Vec::new(),
218                col_field: Vec::new(),
219                vbit_offsets: Vec::new(),
220                repack: Vec::new(),
221            }),
222            // q4tp (§4.10): nibbles from mmap, scale from the row ladder —
223            // 7.3% less file than q4t at the same 4-bit grid.
224            TensorDtype::Q4TiledP if cols % GROUP_SIZE == 0 => Ok(Self::Mapped {
225                model: model.clone(),
226                idx,
227                dtype: entry.dtype,
228                rows,
229                cols,
230                row_scale: Vec::new(),
231                col_field: Vec::new(),
232                vbit_offsets: Vec::new(),
233                repack: Vec::new(),
234            }),
235            // q2tp: 2-bit chunks from mmap, scale from the same row ladder.
236            TensorDtype::Q2TiledP if cols % GROUP_SIZE == 0 => Ok(Self::Mapped {
237                model: model.clone(),
238                idx,
239                dtype: entry.dtype,
240                rows,
241                cols,
242                row_scale: Vec::new(),
243                col_field: Vec::new(),
244                vbit_offsets: Vec::new(),
245                repack: Vec::new(),
246            }),
247            TensorDtype::Q4Block if cols % GROUP_SIZE == 0 => Ok(Self::Mapped {
248                model: model.clone(),
249                idx,
250                dtype: entry.dtype,
251                rows,
252                cols,
253                row_scale: Vec::new(),
254                col_field: Vec::new(),
255                vbit_offsets: Vec::new(),
256                repack: Vec::new(),
257            }),
258            // q1: binary sign-bit tiles from mmap (1-bit-trained models).
259            TensorDtype::Q1 if cols % GROUP_SIZE == 0 => Ok(Self::Mapped {
260                model: model.clone(),
261                idx,
262                dtype: entry.dtype,
263                rows,
264                cols,
265                row_scale: Vec::new(),
266                col_field: Vec::new(),
267                vbit_offsets: Vec::new(),
268                repack: Vec::new(),
269            }),
270            // q1t (ternary + outlier overlay): fused per-row dequant kernel
271            // reads straight from mmap — a 12B q1t stays ~its file size in
272            // RAM instead of dequantizing to ~48 GB of f32.
273            TensorDtype::Q1T if cols % GROUP_SIZE == 0 => Ok(Self::Mapped {
274                model: model.clone(),
275                idx,
276                dtype: entry.dtype,
277                rows,
278                cols,
279                row_scale: Vec::new(),
280                col_field: Vec::new(),
281                vbit_offsets: Vec::new(),
282                repack: Vec::new(),
283            }),
284            // No fused kernel yet → dequantize once (correct, more RAM).
285            _ => {
286                let mut data = vec![0.0f32; rows * cols];
287                cortiq_core::quant::dequant_tensor(entry, bytes, &mut data)?;
288                Ok(Self::from_f32(data, rows, cols))
289            }
290        }
291    }
292
293    /// q1-mapped tensor? (GPU gates: the q1 CPU kernel is
294    /// compute-bound, so offload pays at much smaller shapes than q8.)
295    pub(crate) fn is_q1(&self) -> bool {
296        matches!(
297            self,
298            Self::Mapped {
299                dtype: TensorDtype::Q1,
300                ..
301            }
302        )
303    }
304
305    /// Owned-f32 view (data, rows, cols) — the GDN a/b gate projections
306    /// arrive dequantized (force-f16 in the converter → F32 in RAM).
307    pub(crate) fn f32_parts(&self) -> Option<(&[f32], usize, usize)> {
308        match self {
309            Self::F32 { data, rows, cols } => Some((data, *rows, *cols)),
310            _ => None,
311        }
312    }
313
314    /// (directory idx, rows, cols) of a q1-mapped tensor — the
315    /// whole-block GPU path resolves offsets itself.
316    /// (idx, rows, cols) of a mapped tensor the whole-token GPU graph can drive
317    /// — Q1, Q1T or Q4-block (it resolves the offset and picks the kernel by
318    /// dtype). Q4-block lets a precise down_proj/lm_head stay on-device.
319    /// Named `q1_parts` for historical reasons.
320    pub(crate) fn q1_parts(&self) -> Option<(usize, usize, usize)> {
321        match self {
322            #[cfg(target_os = "macos")]
323            Self::Mapped {
324                dtype: TensorDtype::Q1T,
325                ..
326            } if !crate::gpu::metal_q1t_enabled() => None,
327            Self::Mapped {
328                idx,
329                dtype:
330                    TensorDtype::Q1
331                    | TensorDtype::Q1T
332                    | TensorDtype::Q4Block
333                    | TensorDtype::Q4Tiled
334                    | TensorDtype::Q4TiledP
335                    | TensorDtype::Q2TiledP
336                    | TensorDtype::Q8Row
337                    | TensorDtype::Q8_2f,
338                rows,
339                cols,
340                ..
341            } => Some((*idx, *rows, *cols)),
342            _ => None,
343        }
344    }
345
346    /// (directory idx, rows, cols) of a q4_tiled mapped tensor. The
347    /// chunk-prefill graph takes it in the same 4-tuple slot as
348    /// `q8_row_parts` with an EMPTY row_scale — q4t carries its scales
349    /// inside the 18-byte tiles, and the empty slice is what tells the
350    /// encoder to reach for the q4t kernels.
351    pub(crate) fn q4t_parts(&self) -> Option<(usize, usize, usize)> {
352        match self {
353            Self::Mapped {
354                idx,
355                dtype: TensorDtype::Q4Tiled,
356                rows,
357                cols,
358                ..
359            } => Some((*idx, *rows, *cols)),
360            _ => None,
361        }
362    }
363
364    /// (directory idx, rows, cols) of a q4tp mapped tensor. Same empty-scale
365    /// slot as `q4t_parts` in the chunk graph — the encoder tells the two
366    /// apart by the tensor's dtype, not by the slot.
367    pub(crate) fn q4tp_parts(&self) -> Option<(usize, usize, usize)> {
368        match self {
369            Self::Mapped {
370                idx,
371                dtype: TensorDtype::Q4TiledP,
372                rows,
373                cols,
374                ..
375            } => Some((*idx, *rows, *cols)),
376            _ => None,
377        }
378    }
379
380    /// (directory idx, rows, cols, row_scale) of a plain q8_row mapped
381    /// tensor — the chunk-prefill GPU graph resolves offsets itself.
382    /// q8_2f is excluded on purpose: its column field would need a
383    /// prescale stage on the device.
384    pub(crate) fn q8_row_parts(&self) -> Option<(usize, usize, usize, &[f32])> {
385        match self {
386            Self::Mapped {
387                idx,
388                dtype: TensorDtype::Q8Row,
389                rows,
390                cols,
391                row_scale,
392                col_field,
393                ..
394            } if col_field.is_empty() => Some((*idx, *rows, *cols, row_scale)),
395            _ => None,
396        }
397    }
398
399    /// The layout this tensor is stored in, when it is mapped from a model.
400    /// The frames branch on it — a q2tp gate against a q4tp down is a real
401    /// combination in the 2-bit profile and needs a different kernel.
402    pub fn model_dtype(&self) -> Option<cortiq_core::TensorDtype> {
403        match self {
404            Self::Mapped { dtype, .. } => Some(*dtype),
405            _ => None,
406        }
407    }
408
409    /// The tensor's index in the model directory, when it is mapped from one.
410    /// The GPU frames bind by index rather than by name — a name lookup per
411    /// layer per token is not free, and the index is what the device cache is
412    /// keyed on anyway.
413    pub fn model_idx(&self) -> Option<usize> {
414        match self {
415            Self::Mapped { idx, .. } => Some(*idx),
416            _ => None,
417        }
418    }
419
420    /// The model this tensor is mapped from, when it is mapped at all. The
421    /// GPU frames need the container to reach the bytes; a QTensor already
422    /// holds it, and threading a second handle down every call site to say
423    /// the same thing invites the two to disagree.
424    pub fn model_arc(&self) -> Option<std::sync::Arc<cortiq_core::CmfModel>> {
425        match self {
426            Self::Mapped { model, .. } => Some(model.clone()),
427            _ => None,
428        }
429    }
430
431    pub fn rows(&self) -> usize {
432        match self {
433            Self::F32 { rows, .. } | Self::Mapped { rows, .. } => *rows,
434        }
435    }
436
437    /// Mapped q4t handle (model + directory index) — the fused GPU FFN
438    /// needs the raw file coordinates of its three projections.
439    pub(crate) fn mapped_q4t(&self) -> Option<(&Arc<CmfModel>, usize)> {
440        match self {
441            Self::Mapped {
442                model,
443                idx,
444                dtype: TensorDtype::Q4Tiled,
445                ..
446            } => Some((model, *idx)),
447            _ => None,
448        }
449    }
450
451    /// Same slot as `mapped_q4t` for a q4tp tensor — the fused DiT FFN picks
452    /// its kernels by which of the two answers.
453    pub fn mapped_q4tp(&self) -> Option<(&Arc<CmfModel>, usize)> {
454        match self {
455            Self::Mapped {
456                model,
457                idx,
458                dtype: TensorDtype::Q4TiledP,
459                ..
460            } => Some((model, *idx)),
461            _ => None,
462        }
463    }
464
465    /// (model, tensor idx) for a q2tp mapped weight — the 2-bit twin of
466    /// `mapped_q4tp`, used by the mixed MoE profile.
467    pub fn mapped_q2tp(&self) -> Option<(&Arc<CmfModel>, usize)> {
468        match self {
469            Self::Mapped {
470                model,
471                idx,
472                dtype: TensorDtype::Q2TiledP,
473                ..
474            } => Some((model, *idx)),
475            _ => None,
476        }
477    }
478
479    pub fn cols(&self) -> usize {
480        match self {
481            Self::F32 { cols, .. } | Self::Mapped { cols, .. } => *cols,
482        }
483    }
484
485    /// (model, tensor idx) for a q1 mapped weight — the wgpu token graph
486    /// keys its resident VRAM cache by idx. None for any other dtype/kind.
487    pub fn mapped_q1(&self) -> Option<(&std::sync::Arc<CmfModel>, usize)> {
488        match self {
489            Self::Mapped {
490                model,
491                idx,
492                dtype: TensorDtype::Q1,
493                ..
494            } => Some((model, *idx)),
495            _ => None,
496        }
497    }
498
499    /// (model, idx, kind, row_scale) for a graph-capable mapped weight. kind:
500    /// 0=q8_row (per-row scales), 1=q1, 2=q4_tiled, 3=q1t (tile-embedded, no
501    /// rs). None for dtypes the token graph does not handle (q8_2f/q4_block/vbit).
502    pub fn graph_weight(&self) -> Option<(&std::sync::Arc<CmfModel>, usize, u8, &[f32])> {
503        match self {
504            Self::Mapped {
505                model,
506                idx,
507                dtype: TensorDtype::Q8Row,
508                row_scale,
509                ..
510            } => Some((model, *idx, 0, row_scale.as_slice())),
511            Self::Mapped {
512                model,
513                idx,
514                dtype: TensorDtype::Q1,
515                ..
516            } => Some((model, *idx, 1, &[])),
517            // Q4Tiled is kind 5, NOT 2: both carried 2 historically, and
518            // the wgpu token graph fed 18B interleaved tiles to the
519            // split-layout q4b kernel — garbage output on q4t models
520            // (caught by an end-to-end answer check on real Vulkan).
521            Self::Mapped {
522                model,
523                idx,
524                dtype: TensorDtype::Q4Tiled,
525                ..
526            } => Some((model, *idx, 5, &[])),
527            // Kind 6, not 5: q4tp's nibble stride and scale planes differ,
528            // and feeding them to the q4t kernel is exactly the mistake that
529            // produced garbage when Q4Tiled shared kind 2 with Q4Block.
530            Self::Mapped {
531                model,
532                idx,
533                dtype: TensorDtype::Q4TiledP,
534                ..
535            } => Some((model, *idx, 6, &[])),
536            Self::Mapped {
537                model,
538                idx,
539                dtype: TensorDtype::Q4Block,
540                ..
541            } => Some((model, *idx, 2, &[])),
542            Self::Mapped {
543                model,
544                idx,
545                dtype: TensorDtype::Q1T,
546                ..
547            } => Some((model, *idx, 3, &[])),
548            _ => None,
549        }
550    }
551
552    /// Dense f32 view — only for owned tensors. Masked/sparse execution
553    /// paths require it; quantized weights don't support masks yet.
554    pub fn as_f32(&self) -> Option<&[f32]> {
555        match self {
556            Self::F32 { data, .. } => Some(data),
557            Self::Mapped { .. } => None,
558        }
559    }
560
561    fn quant_bytes(&self) -> &[u8] {
562        match self {
563            Self::Mapped { model, idx, .. } => model.entry_bytes(&model.tensors[*idx]),
564            Self::F32 { .. } => unreachable!("quant_bytes on F32"),
565        }
566    }
567
568    /// Dequantize one row into `dst` (embedding lookup).
569    pub fn row_f32(&self, r: usize, dst: &mut [f32]) {
570        let cols = self.cols();
571        debug_assert_eq!(dst.len(), cols);
572        match self {
573            Self::F32 { data, .. } => dst.copy_from_slice(&data[r * cols..(r + 1) * cols]),
574            Self::Mapped {
575                dtype,
576                row_scale,
577                col_field,
578                vbit_offsets,
579                ..
580            } => {
581                if *dtype == TensorDtype::Q4Tiled {
582                    let bytes = self.quant_bytes();
583                    let gpr = cols / GROUP_SIZE;
584                    for gi in 0..gpr {
585                        let tile = &bytes[(r * gpr + gi) * Q4_TILE..(r * gpr + gi + 1) * Q4_TILE];
586                        let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
587                        for (k, &b) in tile[2..].iter().enumerate() {
588                            dst[gi * GROUP_SIZE + k * 2] = ((b & 0x0F) as f32 - 8.0) * s;
589                            dst[gi * GROUP_SIZE + k * 2 + 1] = (((b >> 4) & 0x0F) as f32 - 8.0) * s;
590                        }
591                    }
592                    return;
593                }
594                if *dtype == TensorDtype::Q4TiledP {
595                    let bytes = self.quant_bytes();
596                    let gpr = cols / GROUP_SIZE;
597                    let v = Q4tpView::new(bytes, self.rows(), cols);
598                    let mut sc = vec![0f32; gpr];
599                    v.scales_into(r, gpr, &mut sc);
600                    for gi in 0..gpr {
601                        let tile = &v.nib[(r * gpr + gi) * Q4TP_NIB..(r * gpr + gi + 1) * Q4TP_NIB];
602                        let s = sc[gi];
603                        for (k, &b) in tile.iter().enumerate() {
604                            dst[gi * GROUP_SIZE + k * 2] = ((b & 0x0F) as f32 - 8.0) * s;
605                            dst[gi * GROUP_SIZE + k * 2 + 1] = (((b >> 4) & 0x0F) as f32 - 8.0) * s;
606                        }
607                    }
608                    return;
609                }
610                if *dtype == TensorDtype::Q2TiledP {
611                    let bytes = self.quant_bytes();
612                    let gpr = cols / GROUP_SIZE;
613                    let v = Q4tpView::new_q2(bytes, self.rows(), cols);
614                    let mut sc = vec![0f32; gpr];
615                    v.scales_into(r, gpr, &mut sc);
616                    for gi in 0..gpr {
617                        let ch =
618                            &v.nib[(r * gpr + gi) * Q2TP_CHUNK..(r * gpr + gi + 1) * Q2TP_CHUNK];
619                        let s = sc[gi];
620                        for (k, &b) in ch.iter().enumerate() {
621                            for j in 0..4 {
622                                dst[gi * GROUP_SIZE + k * 4 + j] =
623                                    (((b >> (2 * j)) & 3) as f32 - 1.5) * s;
624                            }
625                        }
626                    }
627                    return;
628                }
629                if *dtype == TensorDtype::Q4Block {
630                    let (packed, scales) = q4_split(self.quant_bytes(), self.rows(), cols);
631                    let gpr = cols / GROUP_SIZE;
632                    for gi in 0..gpr {
633                        let g = r * gpr + gi;
634                        let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
635                        for (k, &b) in packed[g * 16..(g + 1) * 16].iter().enumerate() {
636                            dst[gi * GROUP_SIZE + k * 2] = ((b & 0x0F) as f32 - 8.0) * s;
637                            dst[gi * GROUP_SIZE + k * 2 + 1] = (((b >> 4) & 0x0F) as f32 - 8.0) * s;
638                        }
639                    }
640                    return;
641                }
642                if *dtype == TensorDtype::Q1 {
643                    let bytes = self.quant_bytes();
644                    let gpr = cols / GROUP_SIZE;
645                    for gi in 0..gpr {
646                        let tile = &bytes[(r * gpr + gi) * Q1_TILE..(r * gpr + gi + 1) * Q1_TILE];
647                        let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
648                        for (j, &b) in tile[2..].iter().enumerate() {
649                            for k in 0..8 {
650                                dst[gi * GROUP_SIZE + j * 8 + k] =
651                                    (((b >> k) & 1) as f32 * 2.0 - 1.0) * s;
652                            }
653                        }
654                    }
655                    return;
656                }
657                if *dtype == TensorDtype::Q1T {
658                    let bytes = self.quant_bytes();
659                    let gpr = cols / GROUP_SIZE;
660                    let base_len = self.rows() * gpr * cortiq_core::quant::Q1T_TILE;
661                    for gi in 0..gpr {
662                        let off = (r * gpr + gi) * cortiq_core::quant::Q1T_TILE;
663                        let s = cortiq_core::quant::f16_to_f32(u16::from_le_bytes([
664                            bytes[off],
665                            bytes[off + 1],
666                        ]));
667                        let codes = &bytes[off + 2..off + cortiq_core::quant::Q1T_TILE];
668                        for k in 0..GROUP_SIZE {
669                            dst[gi * GROUP_SIZE + k] = match cortiq_core::quant::q1t_code(codes, k)
670                            {
671                                1 => s,
672                                2 => -s,
673                                _ => 0.0,
674                            };
675                        }
676                    }
677                    // Overlay
678                    let rows = self.rows();
679                    let entries = base_len + (rows + 1) * 4;
680                    if entries <= bytes.len() {
681                        let ptrs = &bytes[base_len..base_len + (rows + 1) * 4];
682                        let r0 = u32::from_le_bytes([
683                            ptrs[r * 4],
684                            ptrs[r * 4 + 1],
685                            ptrs[r * 4 + 2],
686                            ptrs[r * 4 + 3],
687                        ]) as usize;
688                        let r1 = u32::from_le_bytes([
689                            ptrs[(r + 1) * 4],
690                            ptrs[(r + 1) * 4 + 1],
691                            ptrs[(r + 1) * 4 + 2],
692                            ptrs[(r + 1) * 4 + 3],
693                        ]) as usize;
694                        let off = entries + r0 * 4;
695                        for i in 0..r1 - r0 {
696                            let item = &bytes[off + i * 4..off + i * 4 + 4];
697                            let c = u16::from_le_bytes([item[0], item[1]]) as usize;
698                            let v = cortiq_core::quant::f16_to_f32(u16::from_le_bytes([
699                                item[2], item[3],
700                            ]));
701                            if c < cols {
702                                dst[c] = v;
703                            }
704                        }
705                    }
706                    return;
707                }
708                if matches!(dtype, TensorDtype::Vbit | TensorDtype::VbitRo) {
709                    let bytes = self.quant_bytes();
710                    let rows = self.rows();
711                    let ng = cols / GROUP_SIZE;
712                    let bits = &bytes[..rows];
713                    let sc_off = rows;
714                    // Precomputed at load — embedding lookup used to scan
715                    // the bit-widths of every preceding row (O(token_id)).
716                    let off = vbit_offsets[r];
717                    let b = bits[r] as usize;
718                    let l = ((1usize << (b - 1)) - 1) as f32;
719                    let data = &bytes[off..];
720                    let (mut acc, mut nbits, mut idx) = (0u64, 0usize, 0usize);
721                    for (i, d) in dst.iter_mut().enumerate() {
722                        while nbits < b {
723                            acc = (acc << 8) | data[idx] as u64;
724                            idx += 1;
725                            nbits += 8;
726                        }
727                        let u = ((acc >> (nbits - b)) & ((1u64 << b) - 1)) as f32;
728                        nbits -= b;
729                        let so = (r * ng + i / GROUP_SIZE) * 2;
730                        let sv = f16_to_f32(u16::from_le_bytes([
731                            bytes[sc_off + so],
732                            bytes[sc_off + so + 1],
733                        ]));
734                        *d = (u - l) * sv;
735                    }
736                    return;
737                }
738                let q = &self.quant_bytes()[r * cols..(r + 1) * cols];
739                let s = row_scale[r];
740                match dtype {
741                    TensorDtype::Q8Row => {
742                        for (d, &b) in dst.iter_mut().zip(q) {
743                            *d = (b as i8) as f32 * s;
744                        }
745                    }
746                    TensorDtype::Q8_2f => {
747                        for (i, (d, &b)) in dst.iter_mut().zip(q).enumerate() {
748                            *d = (b as i8) as f32 * s * col_field[i];
749                        }
750                    }
751                    _ => unreachable!(),
752                }
753            }
754        }
755    }
756
757    /// Can this tensor's columns be read cheaply (for sparse down_proj)?
758    /// True for F32/Q8Row/Q8_2f (per-row scale, direct strided access);
759    /// false for group-packed q4/vbit (column access would unpack whole
760    /// groups — sparse execution falls back to f32 for those).
761    pub fn sparse_col_ok(&self) -> bool {
762        match self {
763            Self::F32 { .. } => true,
764            Self::Mapped { dtype, .. } => {
765                matches!(dtype, TensorDtype::Q8Row | TensorDtype::Q8_2f)
766            }
767        }
768    }
769
770    /// down_proj [hidden, inter]: accumulate `w · col(c)` into `out`
771    /// [hidden] — reads ONLY column `c` (one neuron) from the mmap,
772    /// no full-matrix dequant. `out[k] += w · down[k, c]`.
773    pub fn add_col_scaled(&self, c: usize, w: f32, out: &mut [f32]) {
774        let inter = self.cols();
775        let hidden = self.rows();
776        debug_assert_eq!(out.len(), hidden);
777        match self {
778            Self::F32 { data, .. } => {
779                for (k, o) in out.iter_mut().enumerate() {
780                    *o += w * data[k * inter + c];
781                }
782            }
783            Self::Mapped {
784                dtype,
785                row_scale,
786                col_field,
787                ..
788            } => {
789                let q = self.quant_bytes();
790                let colf = if *dtype == TensorDtype::Q8_2f {
791                    col_field[c]
792                } else {
793                    1.0
794                };
795                let wc = w * colf;
796                for (k, o) in out.iter_mut().enumerate() {
797                    let b = q[k * inter + c] as i8 as f32;
798                    *o += wc * b * row_scale[k];
799                }
800            }
801        }
802    }
803
804    /// Dot of row `r` with `x` (gate/up active-neuron path). Reads only
805    /// row `r` from the mmap — no full dequant. q4/vbit dequant the row
806    /// into `scratch` first (rare for active-FFN weights).
807    pub fn row_dot(&self, r: usize, x: &[f32], scratch: &mut [f32]) -> f32 {
808        let cols = self.cols();
809        match self {
810            Self::F32 { data, .. } => {
811                let row = &data[r * cols..(r + 1) * cols];
812                row.iter().zip(x).map(|(w, v)| w * v).sum()
813            }
814            Self::Mapped {
815                dtype,
816                row_scale,
817                col_field,
818                ..
819            } => match dtype {
820                TensorDtype::Q8Row => {
821                    let q = &self.quant_bytes()[r * cols..(r + 1) * cols];
822                    dot_i8_f32(q, x) * row_scale[r]
823                }
824                TensorDtype::Q8_2f => {
825                    let q = &self.quant_bytes()[r * cols..(r + 1) * cols];
826                    dot_i8_col_f32(q, x, col_field) * row_scale[r]
827                }
828                _ => {
829                    self.row_f32(r, scratch);
830                    scratch.iter().zip(x).map(|(w, v)| w * v).sum()
831                }
832            },
833        }
834    }
835
836    /// `out = W · x` (row-major). F32 delegates to the historical
837    /// bit-exact path; Mapped runs the fused int8 kernel.
838    pub fn matvec(&self, x: &[f32], out: &mut [f32], pool: Option<&Pool>) {
839        match self {
840            Self::F32 { data, .. } => matvec_rows(pool, data, x, out),
841            Self::Mapped {
842                model,
843                idx,
844                dtype,
845                rows,
846                cols,
847                row_scale,
848                col_field,
849                vbit_offsets,
850                repack,
851            } => {
852                let _ = (model, idx);
853                if *dtype == TensorDtype::Q4Block {
854                    // GPU route (wgpu q4b kernel) for large q4_block matvecs —
855                    // gives NVIDIA/AMD/Intel q4 models a GPU path. Probe keeps
856                    // the winner; Metal returns false → the CPU kernel below.
857                    if *rows * *cols >= 8_388_608 && crate::gpu::enabled_here() {
858                        let t0 = std::time::Instant::now();
859                        match crate::gpu::probe_arm(crate::gpu::OpClass::Matvec) {
860                            crate::gpu::ProbeArm::Gpu => {
861                                if crate::gpu::q4b_matvec(model, *idx, x, *rows, *cols, out) {
862                                    crate::gpu::probe_record(
863                                        crate::gpu::OpClass::Matvec,
864                                        true,
865                                        t0.elapsed(),
866                                    );
867                                    return;
868                                }
869                            }
870                            crate::gpu::ProbeArm::CpuTimed => {
871                                q4matvec(self.quant_bytes(), x, *rows, *cols, out, pool);
872                                crate::gpu::probe_record(
873                                    crate::gpu::OpClass::Matvec,
874                                    false,
875                                    t0.elapsed(),
876                                );
877                                return;
878                            }
879                            crate::gpu::ProbeArm::Cpu => {}
880                        }
881                    }
882                    q4matvec(self.quant_bytes(), x, *rows, *cols, out, pool);
883                    return;
884                }
885                if *dtype == TensorDtype::Q4Tiled {
886                    q4t_matvec(self.quant_bytes(), x, *rows, *cols, out, pool);
887                    return;
888                }
889                if *dtype == TensorDtype::Q4TiledP {
890                    q4tp_matvec(self.quant_bytes(), x, *rows, *cols, out, pool);
891                    return;
892                }
893                if *dtype == TensorDtype::Q2TiledP {
894                    q2tp_matvec(self.quant_bytes(), x, *rows, *cols, out, pool);
895                    return;
896                }
897                if *dtype == TensorDtype::Q1 {
898                    // GPU route for large q1 matvecs (out_proj / lm_head
899                    // class): the CPU q1 kernel is load-port-bound at
900                    // ~4 GB/s/core, the GPU one is bandwidth-bound — the
901                    // probe measures both arms and keeps the winner.
902                    if *rows * *cols >= 8_388_608 && crate::gpu::enabled_here() {
903                        let t0 = std::time::Instant::now();
904                        let arm = if crate::gpu::q1_force() {
905                            crate::gpu::ProbeArm::Gpu
906                        } else {
907                            crate::gpu::probe_arm(crate::gpu::OpClass::Matvec)
908                        };
909                        match arm {
910                            crate::gpu::ProbeArm::Gpu => {
911                                if crate::gpu::q1_matvec(model, *idx, x, *rows, *cols, out) {
912                                    crate::gpu::probe_record(
913                                        crate::gpu::OpClass::Matvec,
914                                        true,
915                                        t0.elapsed(),
916                                    );
917                                    return;
918                                }
919                            }
920                            crate::gpu::ProbeArm::CpuTimed => {
921                                q1_matvec(self.quant_bytes(), x, *rows, *cols, out, pool);
922                                crate::gpu::probe_record(
923                                    crate::gpu::OpClass::Matvec,
924                                    false,
925                                    t0.elapsed(),
926                                );
927                                return;
928                            }
929                            crate::gpu::ProbeArm::Cpu => {}
930                        }
931                    }
932                    q1_matvec(self.quant_bytes(), x, *rows, *cols, out, pool);
933                    return;
934                }
935                if *dtype == TensorDtype::Q1T {
936                    // GPU route for large q1t matvecs: the ternary BASE dot runs
937                    // on the GPU (load-port-bound on CPU, like q1), then the
938                    // sparse overlay is added on the CPU. Probe keeps the winner.
939                    if *rows * *cols >= 8_388_608 && crate::gpu::enabled_here() {
940                        let t0 = std::time::Instant::now();
941                        match crate::gpu::probe_arm(crate::gpu::OpClass::Matvec) {
942                            crate::gpu::ProbeArm::Gpu => {
943                                if crate::gpu::q1t_matvec(model, *idx, x, *rows, *cols, out) {
944                                    q1t_add_overlay(self.quant_bytes(), x, *rows, *cols, out, pool);
945                                    crate::gpu::probe_record(
946                                        crate::gpu::OpClass::Matvec,
947                                        true,
948                                        t0.elapsed(),
949                                    );
950                                    return;
951                                }
952                            }
953                            crate::gpu::ProbeArm::CpuTimed => {
954                                q1t_matvec(self.quant_bytes(), x, *rows, *cols, out, pool);
955                                crate::gpu::probe_record(
956                                    crate::gpu::OpClass::Matvec,
957                                    false,
958                                    t0.elapsed(),
959                                );
960                                return;
961                            }
962                            crate::gpu::ProbeArm::Cpu => {}
963                        }
964                    }
965                    q1t_matvec(self.quant_bytes(), x, *rows, *cols, out, pool);
966                    return;
967                }
968                if matches!(dtype, TensorDtype::Vbit | TensorDtype::VbitRo) {
969                    vbitmatvec(self.quant_bytes(), vbit_offsets, x, *rows, *cols, out, pool);
970                    return;
971                }
972                let xs = prescale(x, col_field, *dtype);
973                // D5: large q8 matrices (lm_head-class) — hybrid
974                // CPU∥GPU: split the rows, both sides compute
975                // SIMULTANEOUSLY (same math, shared prescale).
976                // GPU share: CMF_GPU_SPLIT (0..1, default 0.5).
977                if *rows >= crate::gpu::min_rows()
978                    && matches!(dtype, TensorDtype::Q8Row | TensorDtype::Q8_2f)
979                    && std::env::var("CMF_GPU_LMHEAD")
980                        .map(|v| v != "0")
981                        .unwrap_or(true)
982                    && crate::gpu::enabled_here()
983                {
984                    // Runtime probe: alternate the hybrid against the
985                    // pure-CPU matvec, keep whichever is faster HERE.
986                    let t0 = std::time::Instant::now();
987                    match crate::gpu::probe_arm(crate::gpu::OpClass::Matvec) {
988                        crate::gpu::ProbeArm::Gpu => {}
989                        crate::gpu::ProbeArm::CpuTimed => {
990                            qmatvec(
991                                self.quant_bytes(),
992                                repack,
993                                row_scale,
994                                x,
995                                col_field,
996                                *dtype,
997                                *rows,
998                                *cols,
999                                out,
1000                                pool,
1001                            );
1002                            crate::gpu::probe_record(
1003                                crate::gpu::OpClass::Matvec,
1004                                false,
1005                                t0.elapsed(),
1006                            );
1007                            return;
1008                        }
1009                        crate::gpu::ProbeArm::Cpu => {
1010                            qmatvec(
1011                                self.quant_bytes(),
1012                                repack,
1013                                row_scale,
1014                                x,
1015                                col_field,
1016                                *dtype,
1017                                *rows,
1018                                *cols,
1019                                out,
1020                                pool,
1021                            );
1022                            return;
1023                        }
1024                    }
1025                    let frac = std::env::var("CMF_GPU_SPLIT")
1026                        .ok()
1027                        .and_then(|v| v.parse::<f32>().ok())
1028                        .unwrap_or(0.5)
1029                        .clamp(0.0, 1.0);
1030                    let cpu_rows = ((*rows as f32) * (1.0 - frac)) as usize;
1031                    let (out_cpu, out_gpu) = out.split_at_mut(cpu_rows);
1032                    let bytes = self.quant_bytes();
1033                    let ok = std::thread::scope(|sc| {
1034                        let g = sc.spawn(|| {
1035                            crate::gpu::q8_matvec_range(
1036                                model,
1037                                *idx,
1038                                cpu_rows,
1039                                &row_scale[cpu_rows..],
1040                                &xs,
1041                                *rows - cpu_rows,
1042                                *cols,
1043                                out_gpu,
1044                            )
1045                        });
1046                        if cpu_rows > 0 {
1047                            // Repack prefix covers the full groups of the
1048                            // CPU half (the split starts at row 0).
1049                            let rep_cpu = if repack.is_empty() {
1050                                &[][..]
1051                            } else {
1052                                &repack[..(cpu_rows / 4) * 4 * *cols]
1053                            };
1054                            qmatvec(
1055                                &bytes[..cpu_rows * *cols],
1056                                rep_cpu,
1057                                &row_scale[..cpu_rows],
1058                                x,
1059                                col_field,
1060                                *dtype,
1061                                cpu_rows,
1062                                *cols,
1063                                out_cpu,
1064                                pool,
1065                            );
1066                        }
1067                        g.join().unwrap_or(false)
1068                    });
1069                    if ok {
1070                        crate::gpu::probe_record(crate::gpu::OpClass::Matvec, true, t0.elapsed());
1071                        return;
1072                    }
1073                    // GPU failed — CPU finishes its half (rows rebased —
1074                    // group offsets don't line up, mmap layout only).
1075                    qmatvec(
1076                        &bytes[cpu_rows * *cols..(*rows) * *cols],
1077                        &[],
1078                        &row_scale[cpu_rows..],
1079                        x,
1080                        col_field,
1081                        *dtype,
1082                        *rows - cpu_rows,
1083                        *cols,
1084                        out_gpu,
1085                        pool,
1086                    );
1087                    return;
1088                }
1089                qmatvec(
1090                    self.quant_bytes(),
1091                    repack,
1092                    row_scale,
1093                    x,
1094                    col_field,
1095                    *dtype,
1096                    *rows,
1097                    *cols,
1098                    out,
1099                    pool,
1100                );
1101            }
1102        }
1103    }
1104
1105    /// Fused two-input matvec (MTP verify pair): weights streamed once.
1106    pub fn matvec2(
1107        &self,
1108        x1: &[f32],
1109        x2: &[f32],
1110        o1: &mut [f32],
1111        o2: &mut [f32],
1112        pool: Option<&Pool>,
1113    ) {
1114        match self {
1115            Self::F32 { data, .. } => matvec_rows2(pool, data, x1, x2, o1, o2),
1116            Self::Mapped {
1117                dtype,
1118                rows,
1119                cols,
1120                row_scale,
1121                col_field,
1122                vbit_offsets,
1123                ..
1124            } => {
1125                if *dtype == TensorDtype::Q4Block {
1126                    q4matvec2(self.quant_bytes(), x1, x2, *rows, *cols, o1, o2, pool);
1127                    return;
1128                }
1129                if *dtype == TensorDtype::Q4Tiled {
1130                    q4t_matvec2(self.quant_bytes(), x1, x2, *rows, *cols, o1, o2, pool);
1131                    return;
1132                }
1133                if *dtype == TensorDtype::Q4TiledP {
1134                    q4tp_matvec2(self.quant_bytes(), x1, x2, *rows, *cols, o1, o2, pool);
1135                    return;
1136                }
1137                if *dtype == TensorDtype::Q2TiledP {
1138                    q2tp_matvec2(self.quant_bytes(), x1, x2, *rows, *cols, o1, o2, pool);
1139                    return;
1140                }
1141                if *dtype == TensorDtype::Q1 {
1142                    q1_matvec2(self.quant_bytes(), x1, x2, *rows, *cols, o1, o2, pool);
1143                    return;
1144                }
1145                if *dtype == TensorDtype::Q1T {
1146                    // Fused ternary pair: one row pass, the register
1147                    // unpack shared across both streams on ARM. (Q1T
1148                    // lacks a row_scale array — scales live inline in
1149                    // the tiles — so it must not fall through to the
1150                    // q8 qmatvec2 below.)
1151                    q1t_matvec2(self.quant_bytes(), x1, x2, *rows, *cols, o1, o2, pool);
1152                    return;
1153                }
1154                if matches!(dtype, TensorDtype::Vbit | TensorDtype::VbitRo) {
1155                    vbitmatvec2(
1156                        self.quant_bytes(),
1157                        vbit_offsets,
1158                        x1,
1159                        x2,
1160                        *rows,
1161                        *cols,
1162                        o1,
1163                        o2,
1164                        pool,
1165                    );
1166                    return;
1167                }
1168                qmatvec2(
1169                    self.quant_bytes(),
1170                    row_scale,
1171                    x1,
1172                    x2,
1173                    col_field,
1174                    *dtype,
1175                    *rows,
1176                    *cols,
1177                    o1,
1178                    o2,
1179                    pool,
1180                );
1181            }
1182        }
1183    }
1184}
1185
1186impl QTensor {
1187    /// Batched matvec (prefill-GEMM): xs — row-major [b, cols],
1188    /// out — row-major [b, rows]. Element-wise semantics are IDENTICAL
1189    /// to b matvec calls (same dot kernels in the same order); the win —
1190    /// the weight row streams from DRAM once per batch, not b times.
1191    pub fn matmat(&self, xs_all: &[f32], b: usize, out: &mut [f32], pool: Option<&Pool>) {
1192        let cols = self.cols();
1193        let rows = self.rows();
1194        debug_assert_eq!(xs_all.len(), b * cols);
1195        debug_assert_eq!(out.len(), b * rows);
1196        // GPTQ calibration: fold this layer's inputs into its Hessian. Only
1197        // Mapped tensors carry a directory name; the check is a relaxed
1198        // atomic load, free when not calibrating.
1199        if crate::gptq_capture::capturing() {
1200            if let Self::Mapped { model, idx, .. } = self {
1201                crate::gptq_capture::accumulate(&model.tensors[*idx].name, xs_all, b, cols);
1202            }
1203        }
1204        match self {
1205            Self::F32 { data, .. } => {
1206                let out_addr = SendMut(out.as_mut_ptr());
1207                let run = |start: usize, end: usize| {
1208                    for o in start..end {
1209                        let row = &data[o * cols..(o + 1) * cols];
1210                        for bi in 0..b {
1211                            let x = &xs_all[bi * cols..(bi + 1) * cols];
1212                            let mut acc = 0f32;
1213                            for j in 0..cols {
1214                                acc += row[j] * x[j];
1215                            }
1216                            unsafe { *out_addr.at(bi * rows + o) = acc };
1217                        }
1218                    }
1219                };
1220                dispatch_rows(pool, rows, &run);
1221            }
1222            Self::Mapped {
1223                dtype,
1224                row_scale,
1225                col_field,
1226                vbit_offsets,
1227                ..
1228            } => {
1229                if *dtype == TensorDtype::Q4Block {
1230                    q4matmat(self.quant_bytes(), xs_all, b, rows, cols, out, pool);
1231                    return;
1232                }
1233                if *dtype == TensorDtype::Q4TiledP {
1234                    // GPU batched q4tp GEMM (dequant + f32nt mul_mm on the
1235                    // device); the probe keeps whichever beats the CPU arm.
1236                    // Narrow (prompt-encode) and wide (DiT) batches probe
1237                    // as separate classes — the regimes have opposite
1238                    // winners and one shared verdict locked the wrong arm.
1239                    // Kill switch (gpu::mm_kill): one grossly slow GPU op
1240                    // (a fair-condition op is ≤~100 ms even at 1024px)
1241                    // means the device is contended by another process
1242                    // (e.g. a simulator) — verdicts are per-process, so
1243                    // without the bail the whole render crawls behind
1244                    // someone else's queue.
1245                    if b >= 32
1246                        && b * rows * cols >= 128_000_000
1247                        && cols % 32 == 0
1248                        && !crate::gpu::mm_killed()
1249                        && crate::gpu::enabled_here()
1250                    {
1251                        let class = if b >= 128 {
1252                            crate::gpu::OpClass::MatmatWide
1253                        } else {
1254                            crate::gpu::OpClass::Matmat
1255                        };
1256                        if let Self::Mapped { model, idx, .. } = self {
1257                            let t0 = std::time::Instant::now();
1258                            match crate::gpu::probe_arm(class) {
1259                                crate::gpu::ProbeArm::Gpu => {
1260                                    if crate::gpu::q4tp_matmat(
1261                                        model, *idx, xs_all, b, rows, cols, out,
1262                                    ) {
1263                                        let el = t0.elapsed();
1264                                        // Work-proportional budget: ~8× the
1265                                        // fair-device estimate (+20 ms slack).
1266                                        // An absolute cap missed the worst
1267                                        // case — contended ops sit at
1268                                        // 100–240 ms each and still bury a
1269                                        // render whose fair op is 3–9 ms.
1270                                        // Cold ops (first PSO build, buffer
1271                                        // alloc) are exempt: a one-off
1272                                        // ~50 ms compile is not contention.
1273                                        let flops = 2.0 * b as f64 * rows as f64 * cols as f64;
1274                                        let budget = std::time::Duration::from_secs_f64(
1275                                            flops / 1.5e12 * 8.0 + 0.020,
1276                                        );
1277                                        if el > budget && !crate::gpu::probe_was_cold() {
1278                                            tracing::warn!(
1279                                                "gpu q4tp matmat took {el:?} (budget {budget:?}) — \
1280                                                 device contended, CPU for the rest of the process"
1281                                            );
1282                                            crate::gpu::mm_kill();
1283                                        }
1284                                        crate::gpu::probe_record(class, true, el);
1285                                        return;
1286                                    }
1287                                }
1288                                crate::gpu::ProbeArm::CpuTimed => {
1289                                    q4tp_matmat(
1290                                        self.quant_bytes(),
1291                                        xs_all,
1292                                        b,
1293                                        rows,
1294                                        cols,
1295                                        out,
1296                                        pool,
1297                                    );
1298                                    crate::gpu::probe_record(class, false, t0.elapsed());
1299                                    return;
1300                                }
1301                                crate::gpu::ProbeArm::Cpu => {}
1302                            }
1303                        }
1304                    }
1305                    q4tp_matmat(self.quant_bytes(), xs_all, b, rows, cols, out, pool);
1306                    return;
1307                }
1308                if *dtype == TensorDtype::Q2TiledP {
1309                    // Without this arm a q2tp tensor falls through to the
1310                    // q8 fallback, which reads it at one BYTE per weight —
1311                    // a 2x overrun that killed pool workers mid-prefill
1312                    // while the dispatcher waited forever.
1313                    q2tp_matmat(self.quant_bytes(), xs_all, b, rows, cols, out, pool);
1314                    return;
1315                }
1316                if *dtype == TensorDtype::Q4Tiled {
1317                    // GPU batched q4t GEMM (dequant + f32nt mul_mm on the
1318                    // device); the probe keeps whichever beats the CPU arm.
1319                    // Narrow (prompt-encode) and wide (DiT) batches probe
1320                    // as separate classes — the regimes have opposite
1321                    // winners and one shared verdict locked the wrong arm.
1322                    // Kill switch (gpu::mm_kill): one grossly slow GPU op
1323                    // (a fair-condition op is ≤~100 ms even at 1024px)
1324                    // means the device is contended by another process
1325                    // (e.g. a simulator) — verdicts are per-process, so
1326                    // without the bail the whole render crawls behind
1327                    // someone else's queue.
1328                    if b >= 32
1329                        && b * rows * cols >= 128_000_000
1330                        && cols % 32 == 0
1331                        && !crate::gpu::mm_killed()
1332                        && crate::gpu::enabled_here()
1333                    {
1334                        let class = if b >= 128 {
1335                            crate::gpu::OpClass::MatmatWide
1336                        } else {
1337                            crate::gpu::OpClass::Matmat
1338                        };
1339                        if let Self::Mapped { model, idx, .. } = self {
1340                            let t0 = std::time::Instant::now();
1341                            match crate::gpu::probe_arm(class) {
1342                                crate::gpu::ProbeArm::Gpu => {
1343                                    if crate::gpu::q4t_matmat(
1344                                        model, *idx, xs_all, b, rows, cols, out,
1345                                    ) {
1346                                        let el = t0.elapsed();
1347                                        // Work-proportional budget: ~8× the
1348                                        // fair-device estimate (+20 ms slack).
1349                                        // An absolute cap missed the worst
1350                                        // case — contended ops sit at
1351                                        // 100–240 ms each and still bury a
1352                                        // render whose fair op is 3–9 ms.
1353                                        // Cold ops (first PSO build, buffer
1354                                        // alloc) are exempt: a one-off
1355                                        // ~50 ms compile is not contention.
1356                                        let flops = 2.0 * b as f64 * rows as f64 * cols as f64;
1357                                        let budget = std::time::Duration::from_secs_f64(
1358                                            flops / 1.5e12 * 8.0 + 0.020,
1359                                        );
1360                                        if el > budget && !crate::gpu::probe_was_cold() {
1361                                            tracing::warn!(
1362                                                "gpu q4t matmat took {el:?} (budget {budget:?}) — \
1363                                                 device contended, CPU for the rest of the process"
1364                                            );
1365                                            crate::gpu::mm_kill();
1366                                        }
1367                                        crate::gpu::probe_record(class, true, el);
1368                                        return;
1369                                    }
1370                                }
1371                                crate::gpu::ProbeArm::CpuTimed => {
1372                                    q4t_matmat(
1373                                        self.quant_bytes(),
1374                                        xs_all,
1375                                        b,
1376                                        rows,
1377                                        cols,
1378                                        out,
1379                                        pool,
1380                                    );
1381                                    crate::gpu::probe_record(class, false, t0.elapsed());
1382                                    return;
1383                                }
1384                                crate::gpu::ProbeArm::Cpu => {}
1385                            }
1386                        }
1387                    }
1388                    q4t_matmat(self.quant_bytes(), xs_all, b, rows, cols, out, pool);
1389                    return;
1390                }
1391                if *dtype == TensorDtype::Q1 {
1392                    // GPU batched q1 GEMM for wide prefill (q1_mul_mm on the
1393                    // device); the probe keeps whichever beats the CPU matmat.
1394                    if b >= 32
1395                        && b * rows * cols >= 128_000_000
1396                        && cols % 64 == 0
1397                        && crate::gpu::enabled_here()
1398                    {
1399                        if let Self::Mapped { model, idx, .. } = self {
1400                            let t0 = std::time::Instant::now();
1401                            match crate::gpu::probe_arm(crate::gpu::OpClass::Matmat) {
1402                                crate::gpu::ProbeArm::Gpu => {
1403                                    if crate::gpu::q1_matmat(
1404                                        model, *idx, xs_all, b, rows, cols, out,
1405                                    ) {
1406                                        crate::gpu::probe_record(
1407                                            crate::gpu::OpClass::Matmat,
1408                                            true,
1409                                            t0.elapsed(),
1410                                        );
1411                                        return;
1412                                    }
1413                                }
1414                                crate::gpu::ProbeArm::CpuTimed => {
1415                                    q1_matmat(self.quant_bytes(), xs_all, b, rows, cols, out, pool);
1416                                    crate::gpu::probe_record(
1417                                        crate::gpu::OpClass::Matmat,
1418                                        false,
1419                                        t0.elapsed(),
1420                                    );
1421                                    return;
1422                                }
1423                                crate::gpu::ProbeArm::Cpu => {}
1424                            }
1425                        }
1426                    }
1427                    q1_matmat(self.quant_bytes(), xs_all, b, rows, cols, out, pool);
1428                    return;
1429                }
1430                if *dtype == TensorDtype::Q1T {
1431                    // GPU batched GEMM for wide prefill (base + overlay on the
1432                    // device); probe keeps the winner vs the CPU matmat.
1433                    if b >= 32 && b * rows * cols >= 128_000_000 && crate::gpu::enabled_here() {
1434                        if let Self::Mapped { model, idx, .. } = self {
1435                            let t0 = std::time::Instant::now();
1436                            match crate::gpu::probe_arm(crate::gpu::OpClass::Matmat) {
1437                                crate::gpu::ProbeArm::Gpu => {
1438                                    if crate::gpu::q1t_matmat(
1439                                        model, *idx, xs_all, b, rows, cols, out,
1440                                    ) {
1441                                        crate::gpu::probe_record(
1442                                            crate::gpu::OpClass::Matmat,
1443                                            true,
1444                                            t0.elapsed(),
1445                                        );
1446                                        return;
1447                                    }
1448                                }
1449                                crate::gpu::ProbeArm::CpuTimed => {
1450                                    q1t_matmat(
1451                                        self.quant_bytes(),
1452                                        xs_all,
1453                                        b,
1454                                        rows,
1455                                        cols,
1456                                        out,
1457                                        pool,
1458                                    );
1459                                    crate::gpu::probe_record(
1460                                        crate::gpu::OpClass::Matmat,
1461                                        false,
1462                                        t0.elapsed(),
1463                                    );
1464                                    return;
1465                                }
1466                                crate::gpu::ProbeArm::Cpu => {}
1467                            }
1468                        }
1469                    }
1470                    q1t_matmat(self.quant_bytes(), xs_all, b, rows, cols, out, pool);
1471                    return;
1472                }
1473                if matches!(dtype, TensorDtype::Vbit | TensorDtype::VbitRo) {
1474                    vbitmatmat(
1475                        self.quant_bytes(),
1476                        vbit_offsets,
1477                        xs_all,
1478                        b,
1479                        rows,
1480                        cols,
1481                        out,
1482                        pool,
1483                    );
1484                    return;
1485                }
1486                let pre: Vec<std::borrow::Cow<'_, [f32]>> = (0..b)
1487                    .map(|bi| prescale(&xs_all[bi * cols..(bi + 1) * cols], col_field, *dtype))
1488                    .collect();
1489                // D5: large prefill-batch GEMMs — on the GPU (threshold by
1490                // work volume: submission carries b×rows×cols MACs).
1491                // Runtime probe: the naive GEMM shader + sync readback
1492                // lose to the CPU GEMM on slow driver stacks — alternate
1493                // both arms and keep the winner.
1494                if b >= 8 && b * rows * cols >= 128_000_000 && crate::gpu::enabled_here() {
1495                    if let Self::Mapped { model, idx, .. } = self {
1496                        let t0 = std::time::Instant::now();
1497                        match crate::gpu::probe_arm(crate::gpu::OpClass::Matmat) {
1498                            crate::gpu::ProbeArm::Gpu
1499                                if crate::gpu::probe_deciding(crate::gpu::OpClass::Matmat)
1500                                    && !crate::gpu::q8_resident_or_upload(model, *idx) =>
1501                            {
1502                                // Cold weights during probing: the upload
1503                                // has started, the count runs on the CPU —
1504                                // the GPU arm samples on the next touch.
1505                                let q = self.quant_bytes();
1506                                qmatmat(q, row_scale, &pre, rows, cols, out, pool);
1507                                return;
1508                            }
1509                            crate::gpu::ProbeArm::Gpu => {
1510                                let flat: Vec<f32> =
1511                                    pre.iter().flat_map(|v| v.iter().copied()).collect();
1512                                if crate::gpu::q8_matmat(
1513                                    model, *idx, row_scale, &flat, b, rows, cols, out,
1514                                ) {
1515                                    crate::gpu::probe_record(
1516                                        crate::gpu::OpClass::Matmat,
1517                                        true,
1518                                        t0.elapsed(),
1519                                    );
1520                                    return;
1521                                }
1522                            }
1523                            crate::gpu::ProbeArm::CpuTimed => {
1524                                let q = self.quant_bytes();
1525                                qmatmat(q, row_scale, &pre, rows, cols, out, pool);
1526                                crate::gpu::probe_record(
1527                                    crate::gpu::OpClass::Matmat,
1528                                    false,
1529                                    t0.elapsed(),
1530                                );
1531                                return;
1532                            }
1533                            crate::gpu::ProbeArm::Cpu => {}
1534                        }
1535                    }
1536                }
1537                let q = self.quant_bytes();
1538                qmatmat(q, row_scale, &pre, rows, cols, out, pool);
1539            }
1540        }
1541    }
1542}
1543
1544impl QTensor {
1545    /// Multi-matrix job (roadmap §3 P0): N tensors sharing one input
1546    /// run under a SINGLE pool dispatch — QKV or gate+up cost one
1547    /// barrier instead of N. Per-row math is the exact same kernel as
1548    /// `matvec` (bit-identical outputs); only the dispatch is fused.
1549    /// Falls back to N sequential matvecs when the set is not a uniform
1550    /// q8-family/F32 group or there is no pool.
1551    pub fn matvec_many<const N: usize>(
1552        ts: [&QTensor; N],
1553        x: &[f32],
1554        mut outs: [&mut [f32]; N],
1555        pool: Option<&Pool>,
1556    ) {
1557        let total_rows: usize = ts.iter().map(|t| t.rows()).sum();
1558        let uniform_q8 = ts.iter().all(|t| {
1559            matches!(
1560                t,
1561                Self::Mapped {
1562                    dtype: TensorDtype::Q8Row | TensorDtype::Q8_2f,
1563                    ..
1564                }
1565            )
1566        });
1567        let uniform_f32 = ts.iter().all(|t| matches!(t, Self::F32 { .. }));
1568        let uniform_q4 = ts.iter().all(|t| {
1569            matches!(
1570                t,
1571                Self::Mapped {
1572                    dtype: TensorDtype::Q4Block,
1573                    ..
1574                }
1575            )
1576        });
1577        let uniform_vbit = ts.iter().all(|t| {
1578            matches!(
1579                t,
1580                Self::Mapped {
1581                    dtype: TensorDtype::Vbit | TensorDtype::VbitRo,
1582                    ..
1583                }
1584            )
1585        });
1586        let uniform_q1 = ts.iter().all(|t| {
1587            matches!(
1588                t,
1589                Self::Mapped {
1590                    dtype: TensorDtype::Q1,
1591                    ..
1592                }
1593            )
1594        });
1595        let uniform_q1t = ts.iter().all(|t| {
1596            matches!(
1597                t,
1598                Self::Mapped {
1599                    dtype: TensorDtype::Q1T,
1600                    ..
1601                }
1602            )
1603        });
1604        let Some(pool) = pool else {
1605            for (t, o) in ts.iter().zip(outs.iter_mut()) {
1606                t.matvec(x, o, None);
1607            }
1608            return;
1609        };
1610        if total_rows < 256
1611            || !(uniform_q8
1612                || uniform_f32
1613                || uniform_q4
1614                || uniform_vbit
1615                || uniform_q1
1616                || uniform_q1t)
1617        {
1618            for (t, o) in ts.iter().zip(outs.iter_mut()) {
1619                t.matvec(x, o, Some(pool));
1620            }
1621            return;
1622        }
1623
1624        if uniform_q1 {
1625            // One shared activation split + group sums (q1 has no col
1626            // field; the same input feeds every tensor).
1627            let outs_addr: [SendMut; N] = std::array::from_fn(|i| SendMut(outs[i].as_mut_ptr()));
1628            if a8w8_enabled() {
1629                let act = split_act(x);
1630                let gsum = q1_group_sums(&act.xq, ts[0].cols() / GROUP_SIZE);
1631                let (act, gsum) = (&act, &gsum);
1632                let closures: [_; N] = std::array::from_fn(|i| {
1633                    let (bytes, gpr, out) =
1634                        (ts[i].quant_bytes(), ts[i].cols() / GROUP_SIZE, outs_addr[i]);
1635                    move |s: usize, e: usize| q1_range_a8w8(bytes, gpr, act, gsum, out, s, e)
1636                });
1637                let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
1638                    std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
1639                pool.run_many(&parts);
1640            } else {
1641                let closures: [_; N] = std::array::from_fn(|i| {
1642                    let (bytes, gpr, out) =
1643                        (ts[i].quant_bytes(), ts[i].cols() / GROUP_SIZE, outs_addr[i]);
1644                    move |s: usize, e: usize| q1_range_f32(bytes, gpr, x, out, s, e)
1645                });
1646                let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
1647                    std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
1648                pool.run_many(&parts);
1649            }
1650            return;
1651        }
1652
1653        if uniform_q1t {
1654            // Q1T batched: one shared activation split + overlay decode,
1655            // all tensors' rows in ONE pool dispatch (saves N−1 dispatches
1656            // and N−1 redundant split_act calls per layer).
1657            let outs_addr: [SendMut; N] = std::array::from_fn(|i| SendMut(outs[i].as_mut_ptr()));
1658            const TILE: usize = cortiq_core::quant::Q1T_TILE;
1659            if a8w8_enabled() {
1660                let act = split_act(x);
1661                let act = &act;
1662                let x_ref = x;
1663                let closures: [_; N] = std::array::from_fn(|i| {
1664                    let bytes = ts[i].quant_bytes();
1665                    let (rows, cols) = (ts[i].rows(), ts[i].cols());
1666                    let gpr = cols / GROUP_SIZE;
1667                    let (rp_off, ent_off, has_ov) = q1t_overlay(bytes, rows * gpr * TILE, rows);
1668                    let out = outs_addr[i];
1669                    move |s: usize, e: usize| {
1670                        q1t_range_a8w8(bytes, gpr, rp_off, ent_off, has_ov, act, x_ref, out, s, e)
1671                    }
1672                });
1673                let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
1674                    std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
1675                pool.run_many(&parts);
1676            } else {
1677                let x_ref = x;
1678                let closures: [_; N] = std::array::from_fn(|i| {
1679                    let bytes = ts[i].quant_bytes();
1680                    let (rows, cols) = (ts[i].rows(), ts[i].cols());
1681                    let gpr = cols / GROUP_SIZE;
1682                    let (rp_off, ent_off, has_ov) = q1t_overlay(bytes, rows * gpr * TILE, rows);
1683                    let out = outs_addr[i];
1684                    move |s: usize, e: usize| {
1685                        q1t_range_f32_batch(bytes, gpr, rp_off, ent_off, has_ov, x_ref, out, s, e)
1686                    }
1687                });
1688                let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
1689                    std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
1690                pool.run_many(&parts);
1691            }
1692            return;
1693        }
1694
1695        if uniform_q4 || uniform_vbit {
1696            let outs_addr: [SendMut; N] = std::array::from_fn(|i| SendMut(outs[i].as_mut_ptr()));
1697            // q4/vbit share one activation split — no per-tensor col field.
1698            if a8w8_enabled() {
1699                let act = split_act(x);
1700                let act = &act;
1701                if uniform_q4 {
1702                    let closures: [_; N] = std::array::from_fn(|i| {
1703                        let (packed, scales) =
1704                            q4_split(ts[i].quant_bytes(), ts[i].rows(), ts[i].cols());
1705                        let (gpr, cols, out) =
1706                            (ts[i].cols() / GROUP_SIZE, ts[i].cols(), outs_addr[i]);
1707                        move |s: usize, e: usize| {
1708                            q4_range_a8w8(packed, scales, gpr, cols, act, out, s, e)
1709                        }
1710                    });
1711                    let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
1712                        std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
1713                    pool.run_many(&parts);
1714                } else {
1715                    let closures: [_; N] = std::array::from_fn(|i| {
1716                        let Self::Mapped { vbit_offsets, .. } = ts[i] else {
1717                            unreachable!()
1718                        };
1719                        let (bytes, rows, cols, out) = (
1720                            ts[i].quant_bytes(),
1721                            ts[i].rows(),
1722                            ts[i].cols(),
1723                            outs_addr[i],
1724                        );
1725                        move |s: usize, e: usize| {
1726                            vbit_range_a8w8(bytes, vbit_offsets, x, act, rows, cols, out, s, e)
1727                        }
1728                    });
1729                    let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
1730                        std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
1731                    pool.run_many(&parts);
1732                }
1733                return;
1734            }
1735            if uniform_q4 {
1736                let closures: [_; N] = std::array::from_fn(|i| {
1737                    let (packed, scales) =
1738                        q4_split(ts[i].quant_bytes(), ts[i].rows(), ts[i].cols());
1739                    let (gpr, out) = (ts[i].cols() / GROUP_SIZE, outs_addr[i]);
1740                    move |s: usize, e: usize| q4_range_f32(packed, scales, gpr, x, out, s, e)
1741                });
1742                let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
1743                    std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
1744                pool.run_many(&parts);
1745            } else {
1746                let closures: [_; N] = std::array::from_fn(|i| {
1747                    let Self::Mapped { vbit_offsets, .. } = ts[i] else {
1748                        unreachable!()
1749                    };
1750                    let (bytes, rows, cols, out) = (
1751                        ts[i].quant_bytes(),
1752                        ts[i].rows(),
1753                        ts[i].cols(),
1754                        outs_addr[i],
1755                    );
1756                    move |s: usize, e: usize| {
1757                        vbit_range_f32(bytes, vbit_offsets, x, rows, cols, out, s, e)
1758                    }
1759                });
1760                let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
1761                    std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
1762                pool.run_many(&parts);
1763            }
1764            return;
1765        }
1766
1767        if uniform_f32 {
1768            let outs_addr: [SendMut; N] = std::array::from_fn(|i| SendMut(outs[i].as_mut_ptr()));
1769            let closures: [_; N] = std::array::from_fn(|i| {
1770                let Self::F32 { data, cols, .. } = ts[i] else {
1771                    unreachable!()
1772                };
1773                let out = outs_addr[i];
1774                move |start: usize, end: usize| {
1775                    for o in start..end {
1776                        let row = &data[o * cols..(o + 1) * cols];
1777                        let mut sum = 0.0f32;
1778                        for j in 0..*cols {
1779                            sum += row[j] * x[j];
1780                        }
1781                        // SAFETY: disjoint (tensor, row) cells per worker.
1782                        unsafe { *out.at(o) = sum };
1783                    }
1784                }
1785            });
1786            let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
1787                std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
1788            pool.run_many(&parts);
1789            return;
1790        }
1791
1792        // Uniform q8-family: per-tensor prescale (q8_2f col fields
1793        // differ per tensor) + the shared range kernels.
1794        struct Ctx<'a> {
1795            bytes: &'a [u8],
1796            #[cfg_attr(not(target_arch = "aarch64"), allow(dead_code))]
1797            rep: &'a [u8],
1798            row_scale: &'a [f32],
1799            cols: usize,
1800            xs: std::borrow::Cow<'a, [f32]>,
1801        }
1802        let ctxs: [Ctx<'_>; N] = std::array::from_fn(|i| {
1803            let Self::Mapped {
1804                dtype,
1805                cols,
1806                row_scale,
1807                col_field,
1808                repack,
1809                ..
1810            } = ts[i]
1811            else {
1812                unreachable!()
1813            };
1814            Ctx {
1815                bytes: ts[i].quant_bytes(),
1816                rep: repack,
1817                row_scale,
1818                cols: *cols,
1819                xs: prescale(x, col_field, *dtype),
1820            }
1821        });
1822        let outs_addr: [SendMut; N] = std::array::from_fn(|i| SendMut(outs[i].as_mut_ptr()));
1823        #[cfg(target_arch = "aarch64")]
1824        if sdot_enabled() {
1825            let acts: [SplitAct; N] = std::array::from_fn(|i| split_act(&ctxs[i].xs));
1826            let closures: [_; N] = std::array::from_fn(|i| {
1827                let (c, act, out) = (&ctxs[i], &acts[i], outs_addr[i]);
1828                move |start: usize, end: usize| {
1829                    q8_range_sdot(c.bytes, c.rep, c.row_scale, act, c.cols, out, start, end)
1830                }
1831            });
1832            let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
1833                std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
1834            pool.run_many(&parts);
1835            return;
1836        }
1837        #[cfg(target_arch = "x86_64")]
1838        if avx2_a8w8_enabled() {
1839            let acts: [SplitAct; N] = std::array::from_fn(|i| split_act(&ctxs[i].xs));
1840            let closures: [_; N] = std::array::from_fn(|i| {
1841                let (c, act, out) = (&ctxs[i], &acts[i], outs_addr[i]);
1842                move |start: usize, end: usize| {
1843                    q8_range_avx2(c.bytes, c.row_scale, act, c.cols, out, start, end)
1844                }
1845            });
1846            let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
1847                std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
1848            pool.run_many(&parts);
1849            return;
1850        }
1851        let closures: [_; N] = std::array::from_fn(|i| {
1852            let (c, out) = (&ctxs[i], outs_addr[i]);
1853            move |start: usize, end: usize| {
1854                q8_range_f32(c.bytes, c.row_scale, &c.xs, c.cols, out, start, end)
1855            }
1856        });
1857        let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
1858            std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
1859        pool.run_many(&parts);
1860    }
1861}
1862
1863impl QTensor {
1864    /// Pair-input multi-matrix job: N tensors × 2 shared inputs under a
1865    /// single pool dispatch — the MTP/pair decode path publishes one job
1866    /// for Q/K/V (and one for gate+up) instead of one per tensor.
1867    /// Per-row math is exactly `matvec2`'s kernels; bit-identical.
1868    #[allow(clippy::needless_range_loop)]
1869    pub fn matvec2_many<const N: usize>(
1870        ts: [&QTensor; N],
1871        x1: &[f32],
1872        x2: &[f32],
1873        mut o1s: [&mut [f32]; N],
1874        mut o2s: [&mut [f32]; N],
1875        pool: Option<&Pool>,
1876    ) {
1877        let total_rows: usize = ts.iter().map(|t| t.rows()).sum();
1878        let uniform_q8 = ts.iter().all(|t| {
1879            matches!(
1880                t,
1881                Self::Mapped {
1882                    dtype: TensorDtype::Q8Row | TensorDtype::Q8_2f,
1883                    ..
1884                }
1885            )
1886        });
1887        let uniform_f32 = ts.iter().all(|t| matches!(t, Self::F32 { .. }));
1888        let uniform_q4 = ts.iter().all(|t| {
1889            matches!(
1890                t,
1891                Self::Mapped {
1892                    dtype: TensorDtype::Q4Block,
1893                    ..
1894                }
1895            )
1896        });
1897        let uniform_vbit = ts.iter().all(|t| {
1898            matches!(
1899                t,
1900                Self::Mapped {
1901                    dtype: TensorDtype::Vbit | TensorDtype::VbitRo,
1902                    ..
1903                }
1904            )
1905        });
1906        let fusable = pool.is_some()
1907            && total_rows >= 256
1908            && (uniform_q8 || uniform_f32 || uniform_q4 || uniform_vbit);
1909        if !fusable {
1910            for i in 0..N {
1911                ts[i].matvec2(x1, x2, o1s[i], o2s[i], pool);
1912            }
1913            return;
1914        }
1915        let pool = pool.unwrap();
1916
1917        if uniform_q4 || uniform_vbit {
1918            let p1: [SendMut; N] = std::array::from_fn(|i| SendMut(o1s[i].as_mut_ptr()));
1919            let p2: [SendMut; N] = std::array::from_fn(|i| SendMut(o2s[i].as_mut_ptr()));
1920            // q4/vbit share activation splits — no per-tensor col field.
1921            if a8w8_enabled() {
1922                let a1 = split_act(x1);
1923                let a2 = split_act(x2);
1924                let (a1, a2) = (&a1, &a2);
1925                if uniform_q4 {
1926                    let closures: [_; N] = std::array::from_fn(|i| {
1927                        let (packed, scales) =
1928                            q4_split(ts[i].quant_bytes(), ts[i].rows(), ts[i].cols());
1929                        let (gpr, cols, o1, o2) =
1930                            (ts[i].cols() / GROUP_SIZE, ts[i].cols(), p1[i], p2[i]);
1931                        move |s: usize, e: usize| {
1932                            q4_range2_a8w8(packed, scales, gpr, cols, a1, a2, o1, o2, s, e)
1933                        }
1934                    });
1935                    let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
1936                        std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
1937                    pool.run_many(&parts);
1938                } else {
1939                    let closures: [_; N] = std::array::from_fn(|i| {
1940                        let Self::Mapped { vbit_offsets, .. } = ts[i] else {
1941                            unreachable!()
1942                        };
1943                        let (bytes, rows, cols, o1, o2) = (
1944                            ts[i].quant_bytes(),
1945                            ts[i].rows(),
1946                            ts[i].cols(),
1947                            p1[i],
1948                            p2[i],
1949                        );
1950                        move |s: usize, e: usize| {
1951                            vbit_range2_a8w8(
1952                                bytes,
1953                                vbit_offsets,
1954                                x1,
1955                                x2,
1956                                a1,
1957                                a2,
1958                                rows,
1959                                cols,
1960                                o1,
1961                                o2,
1962                                s,
1963                                e,
1964                            )
1965                        }
1966                    });
1967                    let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
1968                        std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
1969                    pool.run_many(&parts);
1970                }
1971                return;
1972            }
1973            if uniform_q4 {
1974                let closures: [_; N] = std::array::from_fn(|i| {
1975                    let (packed, scales) =
1976                        q4_split(ts[i].quant_bytes(), ts[i].rows(), ts[i].cols());
1977                    let (gpr, o1, o2) = (ts[i].cols() / GROUP_SIZE, p1[i], p2[i]);
1978                    move |s: usize, e: usize| {
1979                        q4_range2_f32(packed, scales, gpr, x1, x2, o1, o2, s, e)
1980                    }
1981                });
1982                let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
1983                    std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
1984                pool.run_many(&parts);
1985            } else {
1986                let closures: [_; N] = std::array::from_fn(|i| {
1987                    let Self::Mapped { vbit_offsets, .. } = ts[i] else {
1988                        unreachable!()
1989                    };
1990                    let (bytes, rows, cols, o1, o2) = (
1991                        ts[i].quant_bytes(),
1992                        ts[i].rows(),
1993                        ts[i].cols(),
1994                        p1[i],
1995                        p2[i],
1996                    );
1997                    move |s: usize, e: usize| {
1998                        vbit_range2_f32(bytes, vbit_offsets, x1, x2, rows, cols, o1, o2, s, e)
1999                    }
2000                });
2001                let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
2002                    std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
2003                pool.run_many(&parts);
2004            }
2005            return;
2006        }
2007
2008        if uniform_f32 {
2009            let p1: [SendMut; N] = std::array::from_fn(|i| SendMut(o1s[i].as_mut_ptr()));
2010            let p2: [SendMut; N] = std::array::from_fn(|i| SendMut(o2s[i].as_mut_ptr()));
2011            let closures: [_; N] = std::array::from_fn(|i| {
2012                let Self::F32 { data, cols, .. } = ts[i] else {
2013                    unreachable!()
2014                };
2015                let (o1, o2) = (p1[i], p2[i]);
2016                move |start: usize, end: usize| {
2017                    for o in start..end {
2018                        let row = &data[o * cols..(o + 1) * cols];
2019                        let (mut s1, mut s2) = (0.0f32, 0.0f32);
2020                        for j in 0..*cols {
2021                            s1 += row[j] * x1[j];
2022                            s2 += row[j] * x2[j];
2023                        }
2024                        // SAFETY: disjoint (tensor, row) cells per worker.
2025                        unsafe {
2026                            *o1.at(o) = s1;
2027                            *o2.at(o) = s2;
2028                        }
2029                    }
2030                }
2031            });
2032            let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
2033                std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
2034            pool.run_many(&parts);
2035            return;
2036        }
2037
2038        struct Ctx<'a> {
2039            bytes: &'a [u8],
2040            row_scale: &'a [f32],
2041            cols: usize,
2042            xs1: std::borrow::Cow<'a, [f32]>,
2043            xs2: std::borrow::Cow<'a, [f32]>,
2044        }
2045        let ctxs: [Ctx<'_>; N] = std::array::from_fn(|i| {
2046            let Self::Mapped {
2047                dtype,
2048                cols,
2049                row_scale,
2050                col_field,
2051                ..
2052            } = ts[i]
2053            else {
2054                unreachable!()
2055            };
2056            Ctx {
2057                bytes: ts[i].quant_bytes(),
2058                row_scale,
2059                cols: *cols,
2060                xs1: prescale(x1, col_field, *dtype),
2061                xs2: prescale(x2, col_field, *dtype),
2062            }
2063        });
2064        let p1: [SendMut; N] = std::array::from_fn(|i| SendMut(o1s[i].as_mut_ptr()));
2065        let p2: [SendMut; N] = std::array::from_fn(|i| SendMut(o2s[i].as_mut_ptr()));
2066        #[cfg(target_arch = "aarch64")]
2067        if sdot_enabled() {
2068            let acts: [(SplitAct, SplitAct); N] =
2069                std::array::from_fn(|i| (split_act(&ctxs[i].xs1), split_act(&ctxs[i].xs2)));
2070            let closures: [_; N] = std::array::from_fn(|i| {
2071                let (c, a, o1, o2) = (&ctxs[i], &acts[i], p1[i], p2[i]);
2072                move |start: usize, end: usize| {
2073                    q8_range2_sdot(c.bytes, c.row_scale, &a.0, &a.1, c.cols, o1, o2, start, end)
2074                }
2075            });
2076            let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
2077                std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
2078            pool.run_many(&parts);
2079            return;
2080        }
2081        #[cfg(target_arch = "x86_64")]
2082        if avx2_a8w8_enabled() {
2083            let acts: [(SplitAct, SplitAct); N] =
2084                std::array::from_fn(|i| (split_act(&ctxs[i].xs1), split_act(&ctxs[i].xs2)));
2085            let closures: [_; N] = std::array::from_fn(|i| {
2086                let (c, a, o1, o2) = (&ctxs[i], &acts[i], p1[i], p2[i]);
2087                move |start: usize, end: usize| {
2088                    q8_range2_avx2(c.bytes, c.row_scale, &a.0, &a.1, c.cols, o1, o2, start, end)
2089                }
2090            });
2091            let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
2092                std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
2093            pool.run_many(&parts);
2094            return;
2095        }
2096        let closures: [_; N] = std::array::from_fn(|i| {
2097            let (c, o1, o2) = (&ctxs[i], p1[i], p2[i]);
2098            move |start: usize, end: usize| {
2099                q8_range2_f32(
2100                    c.bytes,
2101                    c.row_scale,
2102                    &c.xs1,
2103                    &c.xs2,
2104                    c.cols,
2105                    o1,
2106                    o2,
2107                    start,
2108                    end,
2109                )
2110            }
2111        });
2112        let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
2113            std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
2114        pool.run_many(&parts);
2115    }
2116
2117    /// Fused gate+up matvec with SiLU·mul: for each row r, computes
2118    /// `silu(gate·x) * (up·x)` and writes to `out[r]`. ONE pool dispatch,
2119    /// no intermediate g/u buffers, no separate silu pass. Falls back
2120    /// (returns false) for unsupported dtype combos.
2121    pub fn matvec_silu_mul(
2122        gate: &QTensor,
2123        up: &QTensor,
2124        x: &[f32],
2125        out: &mut [f32],
2126        pool: Option<&Pool>,
2127    ) -> bool {
2128        let inter = gate.rows();
2129        debug_assert_eq!(up.rows(), inter);
2130        debug_assert_eq!(out.len(), inter);
2131        debug_assert_eq!(gate.cols(), up.cols());
2132        if !a8w8_enabled() {
2133            return false;
2134        }
2135        let act = split_act(x);
2136        let act = &act;
2137        let x_ref = x;
2138        let out_addr = SendMut(out.as_mut_ptr());
2139
2140        match (gate, up) {
2141            // Q4Block gate + Q4Block up (most common mobile q4 models)
2142            (
2143                Self::Mapped {
2144                    dtype: TensorDtype::Q4Block,
2145                    ..
2146                },
2147                Self::Mapped {
2148                    dtype: TensorDtype::Q4Block,
2149                    ..
2150                },
2151            ) => {
2152                let (gp, gs) = q4_split(gate.quant_bytes(), gate.rows(), gate.cols());
2153                let (up_p, up_s) = q4_split(up.quant_bytes(), up.rows(), up.cols());
2154                let gpr = gate.cols() / GROUP_SIZE;
2155                let cols = gate.cols();
2156                let run = move |start: usize, end: usize| {
2157                    for r in start..end {
2158                        let mut gv = dot_q4_row_i8(gp, gs, r * gpr, gpr, &act.xq) * act.sx;
2159                        let mut uv = dot_q4_row_i8(up_p, up_s, r * gpr, gpr, &act.xq) * act.sx;
2160                        for &(j, xv) in &act.outliers {
2161                            let flat = r * cols + j;
2162                            let gb = gp[flat / 2];
2163                            let gn = if flat & 1 == 0 { gb & 0x0F } else { gb >> 4 };
2164                            let gsc = f16_to_f32(u16::from_le_bytes([
2165                                gs[(flat / GROUP_SIZE) * 2],
2166                                gs[(flat / GROUP_SIZE) * 2 + 1],
2167                            ]));
2168                            gv += ((gn as i32 - 8) as f32) * gsc * xv;
2169                            let ub = up_p[flat / 2];
2170                            let un = if flat & 1 == 0 { ub & 0x0F } else { ub >> 4 };
2171                            let usc = f16_to_f32(u16::from_le_bytes([
2172                                up_s[(flat / GROUP_SIZE) * 2],
2173                                up_s[(flat / GROUP_SIZE) * 2 + 1],
2174                            ]));
2175                            uv += ((un as i32 - 8) as f32) * usc * xv;
2176                        }
2177                        let silu_g = gv / (1.0 + (-gv).exp());
2178                        // SAFETY: disjoint row ranges per worker.
2179                        unsafe { *out_addr.at(r) = silu_g * uv };
2180                    }
2181                };
2182                dispatch_rows(pool, inter, &run);
2183                true
2184            }
2185            // Q4Tiled gate + Q4Tiled up — one row pass, both tile
2186            // streams sequential, silu·mul fused (same per-row math as
2187            // `q4t_matvec`).
2188            (
2189                Self::Mapped {
2190                    dtype: TensorDtype::Q4Tiled,
2191                    ..
2192                },
2193                Self::Mapped {
2194                    dtype: TensorDtype::Q4Tiled,
2195                    ..
2196                },
2197            ) => {
2198                let g_bytes = gate.quant_bytes();
2199                let u_bytes = up.quant_bytes();
2200                let gpr = gate.cols() / GROUP_SIZE;
2201                let run = move |start: usize, end: usize| {
2202                    for r in start..end {
2203                        let mut gv = dot_q4t_row_i8(g_bytes, r, gpr, &act.xq) * act.sx;
2204                        let mut uv = dot_q4t_row_i8(u_bytes, r, gpr, &act.xq) * act.sx;
2205                        for &(j, xv) in &act.outliers {
2206                            let (w, s) = q4t_outlier(g_bytes, r, gpr, j);
2207                            gv += w * s * xv;
2208                            let (w, s) = q4t_outlier(u_bytes, r, gpr, j);
2209                            uv += w * s * xv;
2210                        }
2211                        let silu_g = gv / (1.0 + (-gv).exp());
2212                        // SAFETY: disjoint row ranges per worker.
2213                        unsafe { *out_addr.at(r) = silu_g * uv };
2214                    }
2215                };
2216                dispatch_rows(pool, inter, &run);
2217                true
2218            }
2219            // Q4TiledP gate + Q4TiledP up — the same fused row pass, with
2220            // each row's two ladders built once and spent on both streams.
2221            (
2222                Self::Mapped {
2223                    dtype: TensorDtype::Q4TiledP,
2224                    ..
2225                },
2226                Self::Mapped {
2227                    dtype: TensorDtype::Q4TiledP,
2228                    ..
2229                },
2230            ) => {
2231                let cols = gate.cols();
2232                let gpr = cols / GROUP_SIZE;
2233                let gv_view = Q4tpView::new(gate.quant_bytes(), inter, cols);
2234                let uv_view = Q4tpView::new(up.quant_bytes(), inter, cols);
2235                let run = |start: usize, end: usize| {
2236                    let (mut gsc, mut usc) = (vec![0f32; gpr], vec![0f32; gpr]);
2237                    for r in start..end {
2238                        gv_view.scales_into(r, gpr, &mut gsc);
2239                        uv_view.scales_into(r, gpr, &mut usc);
2240                        let mut gv = dot_q4tp_row_i8(gv_view.nib, r, gpr, &act.xq, &gsc) * act.sx;
2241                        let mut uv = dot_q4tp_row_i8(uv_view.nib, r, gpr, &act.xq, &usc) * act.sx;
2242                        for &(j, xv) in &act.outliers {
2243                            let (w, s) = q4tp_outlier(gv_view.nib, r, gpr, j, &gsc);
2244                            gv += w * s * xv;
2245                            let (w, s) = q4tp_outlier(uv_view.nib, r, gpr, j, &usc);
2246                            uv += w * s * xv;
2247                        }
2248                        let silu_g = gv / (1.0 + (-gv).exp());
2249                        // SAFETY: disjoint row ranges per worker.
2250                        unsafe { *out_addr.at(r) = silu_g * uv };
2251                    }
2252                };
2253                dispatch_rows(pool, inter, &run);
2254                true
2255            }
2256            // Q1T gate + Q1T up
2257            (
2258                Self::Mapped {
2259                    dtype: TensorDtype::Q1T,
2260                    ..
2261                },
2262                Self::Mapped {
2263                    dtype: TensorDtype::Q1T,
2264                    ..
2265                },
2266            ) => {
2267                const TILE: usize = cortiq_core::quant::Q1T_TILE;
2268                let g_bytes = gate.quant_bytes();
2269                let u_bytes = up.quant_bytes();
2270                let gpr = gate.cols() / GROUP_SIZE;
2271                let (g_rp, g_ent, g_ov) = q1t_overlay(g_bytes, inter * gpr * TILE, inter);
2272                let (u_rp, u_ent, u_ov) = q1t_overlay(u_bytes, inter * gpr * TILE, inter);
2273                let run = move |start: usize, end: usize| {
2274                    for r in start..end {
2275                        let mut gv = q1t_dot_row_i8(g_bytes, r, gpr, &act.xq) * act.sx;
2276                        let mut uv = q1t_dot_row_i8(u_bytes, r, gpr, &act.xq) * act.sx;
2277                        for &(j, xv) in &act.outliers {
2278                            gv += q1t_base_weight(g_bytes, r, gpr, j) * xv;
2279                            uv += q1t_base_weight(u_bytes, r, gpr, j) * xv;
2280                        }
2281                        gv += q1t_row_outlier_correction(g_bytes, r, g_rp, g_ent, g_ov, x_ref);
2282                        uv += q1t_row_outlier_correction(u_bytes, r, u_rp, u_ent, u_ov, x_ref);
2283                        let silu_g = gv / (1.0 + (-gv).exp());
2284                        // SAFETY: disjoint row ranges per worker.
2285                        unsafe { *out_addr.at(r) = silu_g * uv };
2286                    }
2287                };
2288                dispatch_rows(pool, inter, &run);
2289                true
2290            }
2291            _ => false,
2292        }
2293    }
2294
2295    /// Every routed expert's fused gate/up/SiLU under ONE pool dispatch.
2296    ///
2297    /// The per-expert path pays a pool barrier per expert per stage: at 9
2298    /// experts over 40 layers that is ~720 barriers a token, and a decode
2299    /// profile of Qwen3.6-35B-A3B showed the pool parked in
2300    /// `psynch_cvwait` about twice as long as it spent computing. Laying
2301    /// every expert's rows end-to-end in one virtual row space collapses
2302    /// the stage to a single dispatch. The per-row body is the
2303    /// single-expert q4tp arm verbatim, so outputs are bit-identical.
2304    ///
2305    /// `false` = something is outside the fused q4tp kernel (dtype, shape,
2306    /// or the `CMF_SDOT=0` exact contract); the caller walks the ordinary
2307    /// per-expert path.
2308    pub fn moe_gate_up_many(
2309        pairs: &[(&QTensor, &QTensor)],
2310        x: &[f32],
2311        outs: &mut [Vec<f32>],
2312        pool: Option<&Pool>,
2313    ) -> bool {
2314        if pairs.is_empty() || pairs.len() != outs.len() || !a8w8_enabled() {
2315            return false;
2316        }
2317        let inter = pairs[0].0.rows();
2318        let cols = pairs[0].0.cols();
2319        if cols % GROUP_SIZE != 0 {
2320            return false;
2321        }
2322        let gpr = cols / GROUP_SIZE;
2323        let mut views = Vec::with_capacity(pairs.len() * 2);
2324        for ((g, u), o) in pairs.iter().zip(outs.iter()) {
2325            let both_q4tp = matches!(
2326                g,
2327                Self::Mapped {
2328                    dtype: TensorDtype::Q4TiledP,
2329                    ..
2330                }
2331            ) && matches!(
2332                u,
2333                Self::Mapped {
2334                    dtype: TensorDtype::Q4TiledP,
2335                    ..
2336                }
2337            );
2338            if !both_q4tp
2339                || g.rows() != inter
2340                || u.rows() != inter
2341                || g.cols() != cols
2342                || u.cols() != cols
2343                || o.len() != inter
2344            {
2345                return false;
2346            }
2347            views.push(Q4tpView::new(g.quant_bytes(), inter, cols));
2348            views.push(Q4tpView::new(u.quant_bytes(), inter, cols));
2349        }
2350        let act = split_act(x);
2351        let act = &act;
2352        let ptrs: Vec<SendMut> = outs.iter_mut().map(|o| SendMut(o.as_mut_ptr())).collect();
2353        let (views, ptrs) = (&views, &ptrs);
2354        let run = |start: usize, end: usize| {
2355            let (mut gsc, mut usc) = (vec![0f32; gpr], vec![0f32; gpr]);
2356            for flat in start..end {
2357                let (e, r) = (flat / inter, flat % inter);
2358                let gv_view = &views[e * 2];
2359                let uv_view = &views[e * 2 + 1];
2360                gv_view.scales_into(r, gpr, &mut gsc);
2361                uv_view.scales_into(r, gpr, &mut usc);
2362                let mut gv = dot_q4tp_row_i8(gv_view.nib, r, gpr, &act.xq, &gsc) * act.sx;
2363                let mut uv = dot_q4tp_row_i8(uv_view.nib, r, gpr, &act.xq, &usc) * act.sx;
2364                for &(j, xv) in &act.outliers {
2365                    let (w, s) = q4tp_outlier(gv_view.nib, r, gpr, j, &gsc);
2366                    gv += w * s * xv;
2367                    let (w, s) = q4tp_outlier(uv_view.nib, r, gpr, j, &usc);
2368                    uv += w * s * xv;
2369                }
2370                let silu_g = gv / (1.0 + (-gv).exp());
2371                // SAFETY: one worker owns each (expert, row) pair.
2372                unsafe { *ptrs[e].at(r) = silu_g * uv };
2373            }
2374        };
2375        dispatch_rows(pool, pairs.len() * inter, &run);
2376        true
2377    }
2378
2379    /// Every routed expert's down projection, weighted and summed into
2380    /// `out`, under ONE pool dispatch.
2381    ///
2382    /// Partitioned by OUTPUT row rather than by expert: each row is owned
2383    /// by a single worker, so the experts are summed in the caller's order
2384    /// — the same sequence of f32 adds the serial `out[i] += w·eo[i]` loop
2385    /// performs, hence bit-identical. Partitioning by expert instead would
2386    /// race on the shared accumulator.
2387    pub fn moe_down_many(
2388        downs: &[&QTensor],
2389        gs: &[Vec<f32>],
2390        weights: &[f32],
2391        out: &mut [f32],
2392        pool: Option<&Pool>,
2393    ) -> bool {
2394        if downs.is_empty()
2395            || downs.len() != gs.len()
2396            || downs.len() != weights.len()
2397            || !a8w8_enabled()
2398        {
2399            return false;
2400        }
2401        let rows = out.len();
2402        let cols = downs[0].cols();
2403        if cols % GROUP_SIZE != 0 {
2404            return false;
2405        }
2406        let gpr = cols / GROUP_SIZE;
2407        let mut views = Vec::with_capacity(downs.len());
2408        for (d, g) in downs.iter().zip(gs.iter()) {
2409            if !matches!(
2410                d,
2411                Self::Mapped {
2412                    dtype: TensorDtype::Q4TiledP,
2413                    ..
2414                }
2415            ) || d.rows() != rows
2416                || d.cols() != cols
2417                || g.len() != cols
2418            {
2419                return false;
2420            }
2421            views.push(Q4tpView::new(d.quant_bytes(), rows, cols));
2422        }
2423        // One int8 split per expert — the activation vectors differ.
2424        let acts: Vec<SplitAct> = gs.iter().map(|g| split_act(g)).collect();
2425        // Partitioned by OUTPUT row, with the experts folded inside: each
2426        // row is owned by one worker, so they are summed in the caller's
2427        // order — the same f32 sequence the serial `out[i] += w·eo[i]`
2428        // loop produces. Partitioning by expert instead would either race
2429        // on the accumulator or need a scratch plane and a second pass;
2430        // measured, that variant was a wash, so this keeps the simpler
2431        // shape.
2432        let out_addr = SendMut(out.as_mut_ptr());
2433        let (views, acts, weights) = (&views, &acts, &weights);
2434        let run = |start: usize, end: usize| {
2435            let mut sc = vec![0f32; gpr];
2436            for r in start..end {
2437                let mut acc = 0f32;
2438                for (e, v) in views.iter().enumerate() {
2439                    v.scales_into(r, gpr, &mut sc);
2440                    let a = &acts[e];
2441                    let mut d = dot_q4tp_row_i8(v.nib, r, gpr, &a.xq, &sc) * a.sx;
2442                    for &(j, xv) in &a.outliers {
2443                        let (w, s) = q4tp_outlier(v.nib, r, gpr, j, &sc);
2444                        d += w * s * xv;
2445                    }
2446                    acc += weights[e] * d;
2447                }
2448                // SAFETY: disjoint row ranges per worker.
2449                unsafe { *out_addr.at(r) = acc };
2450            }
2451        };
2452        dispatch_rows(pool, rows, &run);
2453        true
2454    }
2455}
2456
2457/// Batched q8 kernel: same math as qmatvec, the row makes a single
2458/// pass from memory for the whole batch.
2459/// Accelerate CBLAS — the Apple AMX matrix units, the same engine
2460/// llama.cpp's `-ngl 0` prefill rides via ggml-blas.
2461#[cfg(target_os = "macos")]
2462mod accel_blas {
2463    #[link(name = "Accelerate", kind = "framework")]
2464    unsafe extern "C" {
2465        pub fn cblas_sgemm(
2466            order: i32,
2467            trans_a: i32,
2468            trans_b: i32,
2469            m: i32,
2470            n: i32,
2471            k: i32,
2472            alpha: f32,
2473            a: *const f32,
2474            lda: i32,
2475            b: *const f32,
2476            ldb: i32,
2477            beta: f32,
2478            c: *mut f32,
2479            ldc: i32,
2480        );
2481    }
2482}
2483
2484#[cfg(target_os = "macos")]
2485pub(crate) fn accel_gemm_enabled() -> bool {
2486    static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
2487    *ON.get_or_init(|| std::env::var("CMF_ACCEL").map(|v| v != "0").unwrap_or(true))
2488}
2489
2490/// Off macOS the "accel" GEMM is the portable NEON micro-kernel below —
2491/// same entry point, so the batched-attention path opens on mobile.
2492#[cfg(all(target_arch = "aarch64", not(target_os = "macos")))]
2493pub(crate) fn accel_gemm_enabled() -> bool {
2494    true
2495}
2496
2497/// Portable NEON f32 GEMM (row-major, optional Bᵀ): a 4×8 fmla
2498/// micro-kernel with A broadcast against B panels — the mobile stand-in
2499/// for Accelerate in the batched causal attention (QKᵀ and P·V). Not a
2500/// BLAS: shapes here are the attention panels (m ≤ heads·chunk,
2501/// k = head_dim or context), and the goal is removing the per-position
2502/// quadratic wall, not peak GEMM.
2503#[cfg(target_arch = "aarch64")]
2504#[allow(clippy::too_many_arguments)]
2505pub(crate) fn neon_gemm_rm(
2506    m: usize,
2507    n: usize,
2508    k: usize,
2509    alpha: f32,
2510    a: &[f32],
2511    lda: usize,
2512    b_mat: &[f32],
2513    ldb: usize,
2514    b_rows_are_n: bool,
2515    c: &mut [f32],
2516    ldc: usize,
2517) {
2518    debug_assert!(a.len() >= (m - 1) * lda + k);
2519    debug_assert!(c.len() >= (m - 1) * ldc + n);
2520    // SAFETY: bounds asserted above; NEON is baseline on aarch64.
2521    unsafe {
2522        use core::arch::aarch64::*;
2523        let mut i = 0usize;
2524        while i < m {
2525            let mi = (m - i).min(4);
2526            let mut j = 0usize;
2527            while j < n {
2528                let nj = (n - j).min(8);
2529                if mi == 4 && nj == 8 {
2530                    let (mut c0a, mut c0b) = (vdupq_n_f32(0.0), vdupq_n_f32(0.0));
2531                    let (mut c1a, mut c1b) = (vdupq_n_f32(0.0), vdupq_n_f32(0.0));
2532                    let (mut c2a, mut c2b) = (vdupq_n_f32(0.0), vdupq_n_f32(0.0));
2533                    let (mut c3a, mut c3b) = (vdupq_n_f32(0.0), vdupq_n_f32(0.0));
2534                    for p in 0..k {
2535                        let (b0, b1) = if b_rows_are_n {
2536                            // B is [n, k]: column p of Bᵀ = element p of
2537                            // eight consecutive B rows — gathered.
2538                            let base = b_mat.as_ptr().add(j * ldb + p);
2539                            let g = |o: usize| *base.add(o * ldb);
2540                            ([g(0), g(1), g(2), g(3)], [g(4), g(5), g(6), g(7)])
2541                        } else {
2542                            let base = b_mat.as_ptr().add(p * ldb + j);
2543                            (
2544                                [*base, *base.add(1), *base.add(2), *base.add(3)],
2545                                [*base.add(4), *base.add(5), *base.add(6), *base.add(7)],
2546                            )
2547                        };
2548                        let bv0 = vld1q_f32(b0.as_ptr());
2549                        let bv1 = vld1q_f32(b1.as_ptr());
2550                        let a0 = vdupq_n_f32(*a.as_ptr().add(i * lda + p));
2551                        let a1 = vdupq_n_f32(*a.as_ptr().add((i + 1) * lda + p));
2552                        let a2 = vdupq_n_f32(*a.as_ptr().add((i + 2) * lda + p));
2553                        let a3 = vdupq_n_f32(*a.as_ptr().add((i + 3) * lda + p));
2554                        c0a = vfmaq_f32(c0a, a0, bv0);
2555                        c0b = vfmaq_f32(c0b, a0, bv1);
2556                        c1a = vfmaq_f32(c1a, a1, bv0);
2557                        c1b = vfmaq_f32(c1b, a1, bv1);
2558                        c2a = vfmaq_f32(c2a, a2, bv0);
2559                        c2b = vfmaq_f32(c2b, a2, bv1);
2560                        c3a = vfmaq_f32(c3a, a3, bv0);
2561                        c3b = vfmaq_f32(c3b, a3, bv1);
2562                    }
2563                    let al = vdupq_n_f32(alpha);
2564                    for (r, (ca, cb)) in [(c0a, c0b), (c1a, c1b), (c2a, c2b), (c3a, c3b)]
2565                        .iter()
2566                        .enumerate()
2567                    {
2568                        let dst = c.as_mut_ptr().add((i + r) * ldc + j);
2569                        vst1q_f32(dst, vmulq_f32(*ca, al));
2570                        vst1q_f32(dst.add(4), vmulq_f32(*cb, al));
2571                    }
2572                } else {
2573                    for r in 0..mi {
2574                        for q in 0..nj {
2575                            let mut acc = 0f32;
2576                            for p in 0..k {
2577                                let bv = if b_rows_are_n {
2578                                    b_mat[(j + q) * ldb + p]
2579                                } else {
2580                                    b_mat[p * ldb + j + q]
2581                                };
2582                                acc += a[(i + r) * lda + p] * bv;
2583                            }
2584                            c[(i + r) * ldc + j + q] = acc * alpha;
2585                        }
2586                    }
2587                }
2588                j += nj;
2589            }
2590            i += mi;
2591        }
2592    }
2593}
2594
2595/// Off-macOS aarch64: the batched attention rides the NEON micro-GEMM.
2596#[cfg(all(target_arch = "aarch64", not(target_os = "macos")))]
2597#[allow(clippy::too_many_arguments)]
2598pub(crate) fn sgemm_rm(
2599    m: usize,
2600    n: usize,
2601    k: usize,
2602    alpha: f32,
2603    a: &[f32],
2604    lda: usize,
2605    b_mat: &[f32],
2606    ldb: usize,
2607    b_rows_are_n: bool,
2608    c: &mut [f32],
2609    ldc: usize,
2610) {
2611    neon_gemm_rm(m, n, k, alpha, a, lda, b_mat, ldb, b_rows_are_n, c, ldc);
2612}
2613
2614/// Row-major f32 GEMM, exposed for offline tools (the AWNP pass builds a
2615/// per-layer projection and applies it to every expert; a naive triple loop
2616/// would turn a two-minute job into half an hour).
2617#[allow(clippy::too_many_arguments)]
2618pub fn sgemm_public(
2619    m: usize,
2620    n: usize,
2621    k: usize,
2622    alpha: f32,
2623    a: &[f32],
2624    lda: usize,
2625    b_mat: &[f32],
2626    ldb: usize,
2627    b_rows_are_n: bool,
2628    c: &mut [f32],
2629    ldc: usize,
2630) {
2631    #[cfg(any(target_os = "macos", target_arch = "aarch64"))]
2632    {
2633        sgemm_rm(m, n, k, alpha, a, lda, b_mat, ldb, b_rows_are_n, c, ldc);
2634    }
2635    // x86 without Accelerate has no sgemm_rm: the specialized paths there are
2636    // quantized kernels, not an f32 GEMM. Only the offline AWNP pass reaches
2637    // this, so correctness matters and throughput does not — a triple loop is
2638    // the honest fallback rather than a reason to make the tool macOS-only.
2639    #[cfg(not(any(target_os = "macos", target_arch = "aarch64")))]
2640    {
2641        for i in 0..m {
2642            for j in 0..n {
2643                let mut acc = 0f32;
2644                for p in 0..k {
2645                    let bv = if b_rows_are_n {
2646                        b_mat[j * ldb + p]
2647                    } else {
2648                        b_mat[p * ldb + j]
2649                    };
2650                    acc += a[i * lda + p] * bv;
2651                }
2652                c[i * ldc + j] = alpha * acc;
2653            }
2654        }
2655    }
2656}
2657
2658/// Row-major f32 GEMM on Accelerate: C[m,n] = alpha·A[m,k] × B(ᵀ).
2659/// `b_rows_are_n` = true multiplies by Bᵀ where B is stored [n, k].
2660#[cfg(target_os = "macos")]
2661#[allow(clippy::too_many_arguments)]
2662pub(crate) fn sgemm_rm(
2663    m: usize,
2664    n: usize,
2665    k: usize,
2666    alpha: f32,
2667    a: &[f32],
2668    lda: usize,
2669    b_mat: &[f32],
2670    ldb: usize,
2671    b_rows_are_n: bool,
2672    c: &mut [f32],
2673    ldc: usize,
2674) {
2675    debug_assert!(a.len() >= (m - 1) * lda + k);
2676    debug_assert!(c.len() >= (m - 1) * ldc + n);
2677    // Test hook: route the attention GEMMs through the portable NEON
2678    // micro-kernel ON APPLE SILICON — how the mobile batched attend is
2679    // measured without a phone in the loop. (Intel macOS has no NEON —
2680    // the hook is a no-op there, Accelerate continues below.)
2681    #[cfg(target_arch = "aarch64")]
2682    if std::env::var("CMF_FORCE_NEON_GEMM")
2683        .map(|v| v == "1")
2684        .unwrap_or(false)
2685    {
2686        return neon_gemm_rm(m, n, k, alpha, a, lda, b_mat, ldb, b_rows_are_n, c, ldc);
2687    }
2688    unsafe {
2689        accel_blas::cblas_sgemm(
2690            101, // RowMajor
2691            111, // NoTrans A
2692            if b_rows_are_n { 112 } else { 111 },
2693            m as i32,
2694            n as i32,
2695            k as i32,
2696            alpha,
2697            a.as_ptr(),
2698            lda as i32,
2699            b_mat.as_ptr(),
2700            ldb as i32,
2701            0.0,
2702            c.as_mut_ptr(),
2703            ldc as i32,
2704        );
2705    }
2706}
2707
2708/// Prefill GEMM through Accelerate (macOS): dequantize q8 rows into
2709/// f32 tiles (scale folded in, pool-parallel) and multiply each tile
2710/// on the AMX with one row-major sgemm. Tiles live in cache, weights
2711/// stream once. Numerics are f32-GEMM (not the int8 dot): prefill
2712/// logits shift within f32 rounding — tolerance-class, like every
2713/// reduction-order change; decode (M=1) never takes this path.
2714#[cfg(target_os = "macos")]
2715fn qmatmat_accel(
2716    q: &[u8],
2717    row_scale: &[f32],
2718    pre: &[std::borrow::Cow<'_, [f32]>],
2719    rows: usize,
2720    cols: usize,
2721    out: &mut [f32],
2722    pool: Option<&Pool>,
2723) {
2724    // NOTE: double-buffering the dequant against the sgemm (a scoped
2725    // thread driving the pool on tile k+1 while the caller multiplies
2726    // tile k) was tried and LOST ~6%: Accelerate's sgemm is itself
2727    // multithreaded, and the dequant workers just steal its cores.
2728    const TR: usize = 2048;
2729    let b = pre.len();
2730    thread_local! {
2731        static XPANEL: std::cell::RefCell<Vec<f32>> = const { std::cell::RefCell::new(Vec::new()) };
2732        static WTILE: std::cell::RefCell<Vec<f32>> = const { std::cell::RefCell::new(Vec::new()) };
2733    }
2734    XPANEL.with(|xp| {
2735        WTILE.with(|wt| {
2736            let mut xpanel = xp.borrow_mut();
2737            xpanel.clear();
2738            for x in pre {
2739                xpanel.extend_from_slice(x);
2740            }
2741            let mut wtile = wt.borrow_mut();
2742            wtile.resize(TR * cols, 0.0);
2743            let mut r0 = 0usize;
2744            while r0 < rows {
2745                let tr = TR.min(rows - r0);
2746                // Dequant the tile (scale folded) — pool-parallel.
2747                let wt_addr = SendMut(wtile.as_mut_ptr());
2748                let run = |start: usize, end: usize| {
2749                    for r in start..end {
2750                        let row = &q[(r0 + r) * cols..(r0 + r + 1) * cols];
2751                        let s = row_scale[r0 + r];
2752                        // SAFETY: workers cover disjoint r ranges.
2753                        let dst =
2754                            unsafe { std::slice::from_raw_parts_mut(wt_addr.at(r * cols), cols) };
2755                        for (d, &v) in dst.iter_mut().zip(row) {
2756                            *d = (v as i8) as f32 * s;
2757                        }
2758                    }
2759                };
2760                dispatch_rows(pool, tr, &run);
2761                // C[b, tr] (at column r0 of out[b, rows]) = X · Wtileᵀ
2762                unsafe {
2763                    accel_blas::cblas_sgemm(
2764                        101, // RowMajor
2765                        111, // NoTrans A
2766                        112, // Trans B
2767                        b as i32,
2768                        tr as i32,
2769                        cols as i32,
2770                        1.0,
2771                        xpanel.as_ptr(),
2772                        cols as i32,
2773                        wtile.as_ptr(),
2774                        cols as i32,
2775                        0.0,
2776                        out.as_mut_ptr().add(r0),
2777                        rows as i32,
2778                    );
2779                }
2780                r0 += tr;
2781            }
2782        })
2783    });
2784}
2785
2786fn qmatmat(
2787    q: &[u8],
2788    row_scale: &[f32],
2789    pre: &[std::borrow::Cow<'_, [f32]>],
2790    rows: usize,
2791    cols: usize,
2792    out: &mut [f32],
2793    pool: Option<&Pool>,
2794) {
2795    let b = pre.len();
2796    debug_assert_eq!(out.len(), b * rows);
2797    // Big prefill batches ride the AMX (roadmap PR3): the row×batch
2798    // SDOT loop below peaks near the CPU's dot throughput, an order
2799    // below the matrix units. Small tensors and tiny test models stay
2800    // on the exact integer path.
2801    #[cfg(target_os = "macos")]
2802    if b >= 8 && rows * cols >= 500_000 && accel_gemm_enabled() {
2803        qmatmat_accel(q, row_scale, pre, rows, cols, out, pool);
2804        return;
2805    }
2806    #[cfg(target_arch = "aarch64")]
2807    if sdot_enabled() {
2808        let acts: Vec<SplitAct> = pre.iter().map(|x| split_act(x)).collect();
2809        let out_addr = SendMut(out.as_mut_ptr());
2810        // Blocked 2×4 (mobile prefill: no AMX to fall back on — this
2811        // path IS the ARM prefill GEMM off Apple silicon).
2812        let blocked_ok = std::env::var("CMF_X86_BLOCKED")
2813            .map(|v| v != "0")
2814            .unwrap_or(true);
2815        let use_i8mm = i8mm_enabled();
2816        if blocked_ok {
2817            let run = |start: usize, end: usize| {
2818                let mut o = start;
2819                while o < end {
2820                    if o + 2 <= end {
2821                        let r0 = &q[o * cols..(o + 1) * cols];
2822                        let r1 = &q[(o + 1) * cols..(o + 2) * cols];
2823                        let mut bi = 0usize;
2824                        while bi + 4 <= acts.len() {
2825                            let xs = [
2826                                acts[bi].xq.as_slice(),
2827                                acts[bi + 1].xq.as_slice(),
2828                                acts[bi + 2].xq.as_slice(),
2829                                acts[bi + 3].xq.as_slice(),
2830                            ];
2831                            let d = if use_i8mm {
2832                                unsafe { dot_i8_smmla_2x4(r0, r1, xs) }
2833                            } else {
2834                                unsafe { dot_i8_sdot_2x4(r0, r1, xs) }
2835                            };
2836                            for (r, row) in [r0, r1].into_iter().enumerate() {
2837                                for k in 0..4 {
2838                                    let act = &acts[bi + k];
2839                                    let mut v = d[r][k] as f32 * act.sx;
2840                                    for &(j, xv) in &act.outliers {
2841                                        v += (row[j] as i8) as f32 * xv;
2842                                    }
2843                                    unsafe {
2844                                        *out_addr.at((bi + k) * rows + o + r) = v * row_scale[o + r]
2845                                    };
2846                                }
2847                            }
2848                            bi += 4;
2849                        }
2850                        while bi < acts.len() {
2851                            for (r, row) in [r0, r1].into_iter().enumerate() {
2852                                let v = row_dot_sdot(row, &acts[bi]) * row_scale[o + r];
2853                                unsafe { *out_addr.at(bi * rows + o + r) = v };
2854                            }
2855                            bi += 1;
2856                        }
2857                        o += 2;
2858                    } else {
2859                        let row = &q[o * cols..(o + 1) * cols];
2860                        for (bi, act) in acts.iter().enumerate() {
2861                            let v = row_dot_sdot(row, act) * row_scale[o];
2862                            unsafe { *out_addr.at(bi * rows + o) = v };
2863                        }
2864                        o += 1;
2865                    }
2866                }
2867            };
2868            dispatch_rows(pool, rows, &run);
2869            return;
2870        }
2871        let run = |start: usize, end: usize| {
2872            for o in start..end {
2873                let row = &q[o * cols..(o + 1) * cols];
2874                for (bi, act) in acts.iter().enumerate() {
2875                    let v = row_dot_sdot(row, act) * row_scale[o];
2876                    unsafe { *out_addr.at(bi * rows + o) = v };
2877                }
2878            }
2879        };
2880        dispatch_rows(pool, rows, &run);
2881        return;
2882    }
2883    // x86 A8W8 batch. Non-VNNI parts take the BLOCKED 2×4 kernel
2884    // (roadmap P0: two weight rows' abs() stay in registers across four
2885    // activation streams); VNNI machines keep the per-row bias-trick
2886    // dot, which is already throughput-bound there.
2887    #[cfg(target_arch = "x86_64")]
2888    if avx2_a8w8_enabled() {
2889        let acts: Vec<SplitAct> = pre.iter().map(|x| split_act(x)).collect();
2890        let out_addr = SendMut(out.as_mut_ptr());
2891        // CMF_X86_BLOCKED=0 forces the per-row path (paired in-process
2892        // A/B on noisy shared-vCPU hosts).
2893        let blocked_ok = std::env::var("CMF_X86_BLOCKED")
2894            .map(|v| v != "0")
2895            .unwrap_or(true);
2896        if !avx512vnni_enabled() && blocked_ok {
2897            let run = |start: usize, end: usize| {
2898                let mut o = start;
2899                while o < end {
2900                    if o + 2 <= end {
2901                        let r0 = &q[o * cols..(o + 1) * cols];
2902                        let r1 = &q[(o + 1) * cols..(o + 2) * cols];
2903                        let mut bi = 0usize;
2904                        while bi + 4 <= acts.len() {
2905                            let xs = [
2906                                acts[bi].xq.as_slice(),
2907                                acts[bi + 1].xq.as_slice(),
2908                                acts[bi + 2].xq.as_slice(),
2909                                acts[bi + 3].xq.as_slice(),
2910                            ];
2911                            let d = unsafe { dot_i8_i8_avx2_2x4(r0, r1, xs) };
2912                            for (r, row) in [r0, r1].into_iter().enumerate() {
2913                                for k in 0..4 {
2914                                    let act = &acts[bi + k];
2915                                    let mut v = d[r][k] as f32 * act.sx;
2916                                    for &(j, xv) in &act.outliers {
2917                                        v += (row[j] as i8) as f32 * xv;
2918                                    }
2919                                    unsafe {
2920                                        *out_addr.at((bi + k) * rows + o + r) = v * row_scale[o + r]
2921                                    };
2922                                }
2923                            }
2924                            bi += 4;
2925                        }
2926                        while bi < acts.len() {
2927                            for (r, row) in [r0, r1].into_iter().enumerate() {
2928                                let v = row_dot_avx2(row, &acts[bi]) * row_scale[o + r];
2929                                unsafe { *out_addr.at(bi * rows + o + r) = v };
2930                            }
2931                            bi += 1;
2932                        }
2933                        o += 2;
2934                    } else {
2935                        let row = &q[o * cols..(o + 1) * cols];
2936                        for (bi, act) in acts.iter().enumerate() {
2937                            let v = row_dot_avx2(row, act) * row_scale[o];
2938                            unsafe { *out_addr.at(bi * rows + o) = v };
2939                        }
2940                        o += 1;
2941                    }
2942                }
2943            };
2944            dispatch_rows(pool, rows, &run);
2945            return;
2946        }
2947        let run = |start: usize, end: usize| {
2948            for o in start..end {
2949                let row = &q[o * cols..(o + 1) * cols];
2950                for (bi, act) in acts.iter().enumerate() {
2951                    let v = row_dot_avx2(row, act) * row_scale[o];
2952                    unsafe { *out_addr.at(bi * rows + o) = v };
2953                }
2954            }
2955        };
2956        dispatch_rows(pool, rows, &run);
2957        return;
2958    }
2959    let out_addr = SendMut(out.as_mut_ptr());
2960    let run = |start: usize, end: usize| {
2961        for o in start..end {
2962            let row = &q[o * cols..(o + 1) * cols];
2963            for (bi, x) in pre.iter().enumerate() {
2964                let mut acc = 0f32;
2965                for j in 0..cols {
2966                    acc += (row[j] as i8) as f32 * x[j];
2967                }
2968                unsafe { *out_addr.at(bi * rows + o) = acc * row_scale[o] };
2969            }
2970        }
2971    };
2972    dispatch_rows(pool, rows, &run);
2973}
2974
2975/// Split rows across pool workers (shared qmatvec pattern). Self-balancing
2976/// — see `Pool::run_rows` for why a static 1/n split is wrong here.
2977fn dispatch_rows(pool: Option<&Pool>, rows: usize, run: &(dyn Fn(usize, usize) + Sync)) {
2978    match pool {
2979        Some(pool) if rows >= 256 => pool.run_rows(rows, run),
2980        _ => run(0, rows),
2981    }
2982}
2983
2984/// Split a q4_block blob into (packed nibbles, f16 group scales).
2985fn q4_split(bytes: &[u8], rows: usize, cols: usize) -> (&[u8], &[u8]) {
2986    let groups = rows * cols / GROUP_SIZE;
2987    bytes.split_at(groups * 16)
2988}
2989
2990/// SIMD unpack for the dominant vbit width B=4 (94% of rows on the
2991/// log2-shape calibration): 16 packed bytes -> 32 centered i8 values.
2992/// vbit packs MSB-first, so the HIGH nibble is the even element
2993/// (opposite of q4_block's lo-first interleave). Centering is u-7.
2994#[inline]
2995fn vbit_fill4(data: &[u8], buf: &mut [u8]) {
2996    #[cfg(target_arch = "aarch64")]
2997    unsafe {
2998        return vbit_fill4_neon(data, buf);
2999    }
3000    #[cfg(target_arch = "x86_64")]
3001    if avx2_enabled() {
3002        return unsafe { vbit_fill4_avx2(data, buf) };
3003    }
3004    #[allow(unreachable_code)]
3005    for (blk, chunk) in buf.chunks_exact_mut(8).enumerate() {
3006        let u = unpack8::<4>(&data[blk * 4..]);
3007        for k in 0..8 {
3008            chunk[k] = (u[k] - 7) as i8 as u8;
3009        }
3010    }
3011}
3012
3013#[cfg(target_arch = "aarch64")]
3014#[target_feature(enable = "neon")]
3015unsafe fn vbit_fill4_neon(data: &[u8], buf: &mut [u8]) {
3016    // SAFETY: buf.len() is a multiple of GROUP_SIZE=32; data holds
3017    // buf.len()/2 packed bytes (validated at load).
3018    unsafe {
3019        use core::arch::aarch64::*;
3020        let n = buf.len();
3021        let mask = vdupq_n_u8(0x0F);
3022        let seven = vdupq_n_s8(7);
3023        let mut g = 0usize;
3024        while g * 32 + 32 <= n {
3025            let b = vld1q_u8(data.as_ptr().add(g * 16));
3026            let hi = vshrq_n_u8::<4>(b);
3027            let lo = vandq_u8(b, mask);
3028            let z0 = vsubq_s8(vreinterpretq_s8_u8(vzip1q_u8(hi, lo)), seven);
3029            let z1 = vsubq_s8(vreinterpretq_s8_u8(vzip2q_u8(hi, lo)), seven);
3030            vst1q_u8(buf.as_mut_ptr().add(g * 32), vreinterpretq_u8_s8(z0));
3031            vst1q_u8(buf.as_mut_ptr().add(g * 32 + 16), vreinterpretq_u8_s8(z1));
3032            g += 1;
3033        }
3034    }
3035}
3036
3037#[cfg(target_arch = "x86_64")]
3038#[target_feature(enable = "avx2")]
3039unsafe fn vbit_fill4_avx2(data: &[u8], buf: &mut [u8]) {
3040    // SAFETY: see vbit_fill4_neon.
3041    unsafe {
3042        use core::arch::x86_64::*;
3043        let n = buf.len();
3044        let mask = _mm_set1_epi8(0x0F);
3045        let seven = _mm256_set1_epi8(7);
3046        let mut g = 0usize;
3047        while g * 32 + 32 <= n {
3048            let b = _mm_loadu_si128(data.as_ptr().add(g * 16) as *const __m128i);
3049            let hi = _mm_and_si128(_mm_srli_epi16::<4>(b), mask);
3050            let lo = _mm_and_si128(b, mask);
3051            let z = _mm256_sub_epi8(
3052                _mm256_set_m128i(_mm_unpackhi_epi8(hi, lo), _mm_unpacklo_epi8(hi, lo)),
3053                seven,
3054            );
3055            _mm256_storeu_si256(buf.as_mut_ptr().add(g * 32) as *mut __m256i, z);
3056            g += 1;
3057        }
3058    }
3059}
3060
3061/// Unpack 8 MSB-first B-bit values from exactly B bytes (fixed shifts —
3062/// no serial bit-buffer, auto-vectorizable). Every 32-value group starts
3063/// byte-aligned (32·B/8 is integral for B∈3..8), so groups decompose
3064/// into 4 such blocks.
3065#[inline(always)]
3066fn unpack8<const B: usize>(data: &[u8]) -> [i32; 8] {
3067    let mut acc = 0u64;
3068    for i in 0..B {
3069        acc = (acc << 8) | data[i] as u64;
3070    }
3071    let mask = (1u64 << B) - 1;
3072    let mut out = [0i32; 8];
3073    for (k, o) in out.iter_mut().enumerate() {
3074        *o = ((acc >> ((7 - k) * B)) & mask) as i32;
3075    }
3076    out
3077}
3078
3079/// Fused vbit matvec straight from the mapped bytes (spec §3, P13
3080/// FIG.3): [u8 bits: rows][f16 scales: rows·cols/32][bit-packed rows,
3081/// MSB-first, byte-padded]. Row data offsets are precomputed at load
3082/// (`vbit_row_offsets`) — the per-call prefix scan was O(rows) pure
3083/// overhead on every matvec.
3084#[allow(clippy::too_many_arguments)]
3085fn vbitmatvec(
3086    bytes: &[u8],
3087    offsets: &[usize],
3088    x: &[f32],
3089    rows: usize,
3090    cols: usize,
3091    out: &mut [f32],
3092    pool: Option<&Pool>,
3093) {
3094    debug_assert_eq!(out.len(), rows);
3095    debug_assert_eq!(offsets.len(), rows + 1);
3096
3097    // SDOT path: unpack the row to centered i8 once, then per-group
3098    // int8 dot against the quantized activations — same A8W8 contract
3099    // as q8 (bounded noise; CMF_SDOT=0 keeps the exact scalar path).
3100    if a8w8_enabled() {
3101        let act = split_act(x);
3102        let out_addr = SendMut(out.as_mut_ptr());
3103        let run = move |start: usize, end: usize| {
3104            vbit_range_a8w8(bytes, offsets, x, &act, rows, cols, out_addr, start, end)
3105        };
3106        dispatch_rows(pool, rows, &run);
3107        return;
3108    }
3109
3110    let out_addr = SendMut(out.as_mut_ptr());
3111    let run = move |start: usize, end: usize| {
3112        vbit_range_f32(bytes, offsets, x, rows, cols, out_addr, start, end)
3113    };
3114    dispatch_rows(pool, rows, &run);
3115}
3116
3117/// One vbit row range via the A8W8 int8 path — kernel body of
3118/// `vbitmatvec`, extracted so multi-matrix jobs can drive it for
3119/// several tensors in one dispatch (b=8 rows go exact f32).
3120#[allow(clippy::too_many_arguments)]
3121fn vbit_range_a8w8(
3122    bytes: &[u8],
3123    offsets: &[usize],
3124    x: &[f32],
3125    act: &SplitAct,
3126    rows: usize,
3127    cols: usize,
3128    out: SendMut,
3129    start: usize,
3130    end: usize,
3131) {
3132    let ng = cols / GROUP_SIZE;
3133    let bits = &bytes[..rows];
3134    let sc_off = rows;
3135    let row_dot = |r: usize| -> f32 {
3136        let b = bits[r] as usize;
3137        let l = (1i32 << (b - 1)) - 1;
3138        let mask = (1u64 << b) - 1;
3139        let data = &bytes[offsets[r]..offsets[r + 1]];
3140        if b == 8 {
3141            // u−L reaches 128 → does not fit i8; exact f32 path.
3142            let (mut acc, mut nbits, mut idx) = (0u64, 0usize, 0usize);
3143            let mut dot = 0f32;
3144            for g in 0..ng {
3145                let so = (r * ng + g) * 2;
3146                let sgf = f16_to_f32(u16::from_le_bytes([
3147                    bytes[sc_off + so],
3148                    bytes[sc_off + so + 1],
3149                ]));
3150                let xg = &x[g * GROUP_SIZE..(g + 1) * GROUP_SIZE];
3151                let mut gd = 0f32;
3152                for &xv in xg.iter() {
3153                    if nbits < 8 {
3154                        acc = (acc << 8) | data[idx] as u64;
3155                        idx += 1;
3156                        nbits += 8;
3157                    }
3158                    let u = ((acc >> (nbits - 8)) & 0xFF) as i32;
3159                    nbits -= 8;
3160                    gd += (u - l) as f32 * xv;
3161                }
3162                dot += gd * sgf;
3163            }
3164            return dot;
3165        }
3166        // Per-worker scratch: this closure runs for every row of the
3167        // tensor (lm_head ≈ 150k rows/token) — a heap allocation per
3168        // row was measurable pure overhead.
3169        thread_local! {
3170            static VBIT_SCRATCH: std::cell::RefCell<Vec<u8>> =
3171                const { std::cell::RefCell::new(Vec::new()) };
3172        }
3173        #[inline(always)]
3174        fn fill<const B: usize>(data: &[u8], l: i32, buf: &mut [u8]) {
3175            for (blk, chunk) in buf.chunks_exact_mut(8).enumerate() {
3176                let u = unpack8::<B>(&data[blk * B..]);
3177                for k in 0..8 {
3178                    chunk[k] = (u[k] - l) as i8 as u8;
3179                }
3180            }
3181        }
3182        let _ = mask;
3183        VBIT_SCRATCH.with(|scratch| {
3184            let mut buf = scratch.borrow_mut();
3185            buf.resize(cols, 0);
3186            match b {
3187                3 => fill::<3>(data, l, &mut buf),
3188                4 => vbit_fill4(data, &mut buf),
3189                5 => fill::<5>(data, l, &mut buf),
3190                6 => fill::<6>(data, l, &mut buf),
3191                _ => unreachable!(),
3192            }
3193            let mut dot = 0f32;
3194            for g in 0..ng {
3195                let so = (r * ng + g) * 2;
3196                let s = f16_to_f32(u16::from_le_bytes([
3197                    bytes[sc_off + so],
3198                    bytes[sc_off + so + 1],
3199                ]));
3200                let d = dot_i8_i8(
3201                    &buf[g * GROUP_SIZE..(g + 1) * GROUP_SIZE],
3202                    &act.xq[g * GROUP_SIZE..(g + 1) * GROUP_SIZE],
3203                ) as f32
3204                    * act.sx;
3205                dot += d * s;
3206            }
3207            for &(j, xv) in &act.outliers {
3208                let so = (r * ng + j / GROUP_SIZE) * 2;
3209                let s = f16_to_f32(u16::from_le_bytes([
3210                    bytes[sc_off + so],
3211                    bytes[sc_off + so + 1],
3212                ]));
3213                // xq is zeroed at outlier slots — add the exact term.
3214                dot += (buf[j] as i8) as f32 * s * xv;
3215            }
3216            dot
3217        })
3218    };
3219    for r in start..end {
3220        // SAFETY: disjoint row ranges per worker.
3221        unsafe { *out.at(r) = row_dot(r) };
3222    }
3223}
3224
3225/// Exact scalar vbit row range (same extraction, non-SDOT path).
3226#[allow(clippy::too_many_arguments)]
3227fn vbit_range_f32(
3228    bytes: &[u8],
3229    offsets: &[usize],
3230    x: &[f32],
3231    rows: usize,
3232    cols: usize,
3233    out: SendMut,
3234    start: usize,
3235    end: usize,
3236) {
3237    let ng = cols / GROUP_SIZE;
3238    let bits = &bytes[..rows];
3239    let sc_off = rows;
3240    // Per-bit-width specialized inner loops: the compiler unrolls the
3241    // constant shifts (the generic bit-buffer loop was branch-bound —
3242    // 5.6 vs 13.2 tok/s q4 on the 0.8B).
3243    #[inline(always)]
3244    fn dot_row<const B: usize>(
3245        data: &[u8],
3246        bytes: &[u8],
3247        sc_off: usize,
3248        r: usize,
3249        ng: usize,
3250        x: &[f32],
3251    ) -> f32 {
3252        let l = ((1i32 << (B - 1)) - 1) as f32;
3253        let gbytes = GROUP_SIZE * B / 8;
3254        let mut dot = 0f32;
3255        for g in 0..ng {
3256            let so = (r * ng + g) * 2;
3257            let s = f16_to_f32(u16::from_le_bytes([
3258                bytes[sc_off + so],
3259                bytes[sc_off + so + 1],
3260            ]));
3261            let xg = &x[g * GROUP_SIZE..(g + 1) * GROUP_SIZE];
3262            let gd0 = &data[g * gbytes..(g + 1) * gbytes];
3263            let mut gd = 0f32;
3264            for blk in 0..GROUP_SIZE / 8 {
3265                let u = unpack8::<B>(&gd0[blk * B..]);
3266                let xb = &xg[blk * 8..blk * 8 + 8];
3267                for k in 0..8 {
3268                    gd += (u[k] as f32 - l) * xb[k];
3269                }
3270            }
3271            dot += gd * s;
3272        }
3273        dot
3274    }
3275    for r in start..end {
3276        let data = &bytes[offsets[r]..offsets[r + 1]];
3277        let v = match bits[r] {
3278            3 => dot_row::<3>(data, bytes, sc_off, r, ng, x),
3279            4 => dot_row::<4>(data, bytes, sc_off, r, ng, x),
3280            5 => dot_row::<5>(data, bytes, sc_off, r, ng, x),
3281            6 => dot_row::<6>(data, bytes, sc_off, r, ng, x),
3282            8 => dot_row::<8>(data, bytes, sc_off, r, ng, x),
3283            b => unreachable!("vbit bit-width {b} (validated at load)"),
3284        };
3285        // SAFETY: disjoint row ranges per worker.
3286        unsafe { *out.at(r) = v };
3287    }
3288}
3289
3290/// Fused two-input vbit matvec: each row is unpacked from the mmap ONCE
3291/// and dotted against BOTH activations (MTP verify / pair prefill used
3292/// to run two full matvecs — double weight traffic and double unpack).
3293/// Per-input math is identical to `vbitmatvec` → same accuracy contract.
3294#[allow(clippy::too_many_arguments)]
3295fn vbitmatvec2(
3296    bytes: &[u8],
3297    offsets: &[usize],
3298    x1: &[f32],
3299    x2: &[f32],
3300    rows: usize,
3301    cols: usize,
3302    o1: &mut [f32],
3303    o2: &mut [f32],
3304    pool: Option<&Pool>,
3305) {
3306    debug_assert_eq!(o1.len(), rows);
3307    debug_assert_eq!(o2.len(), rows);
3308
3309    if a8w8_enabled() {
3310        let a1 = split_act(x1);
3311        let a2 = split_act(x2);
3312        let p1 = SendMut(o1.as_mut_ptr());
3313        let p2 = SendMut(o2.as_mut_ptr());
3314        let run = move |start: usize, end: usize| {
3315            vbit_range2_a8w8(
3316                bytes, offsets, x1, x2, &a1, &a2, rows, cols, p1, p2, start, end,
3317            )
3318        };
3319        dispatch_rows(pool, rows, &run);
3320        return;
3321    }
3322
3323    let p1 = SendMut(o1.as_mut_ptr());
3324    let p2 = SendMut(o2.as_mut_ptr());
3325    let run = move |start: usize, end: usize| {
3326        vbit_range2_f32(bytes, offsets, x1, x2, rows, cols, p1, p2, start, end)
3327    };
3328    dispatch_rows(pool, rows, &run);
3329}
3330
3331/// Two-input vbit row range via the A8W8 int8 path — kernel body of
3332/// `vbitmatvec2`, extracted for pair multi-matrix jobs (b=8 rows go
3333/// exact f32 for both lanes, bits streamed once).
3334#[allow(clippy::too_many_arguments)]
3335fn vbit_range2_a8w8(
3336    bytes: &[u8],
3337    offsets: &[usize],
3338    x1: &[f32],
3339    x2: &[f32],
3340    a1: &SplitAct,
3341    a2: &SplitAct,
3342    rows: usize,
3343    cols: usize,
3344    p1: SendMut,
3345    p2: SendMut,
3346    start: usize,
3347    end: usize,
3348) {
3349    let ng = cols / GROUP_SIZE;
3350    let bits = &bytes[..rows];
3351    let sc_off = rows;
3352    let row_dots = |r: usize| -> (f32, f32) {
3353        let b = bits[r] as usize;
3354        let l = (1i32 << (b - 1)) - 1;
3355        let data = &bytes[offsets[r]..offsets[r + 1]];
3356        if b == 8 {
3357            // u−L reaches 128 → does not fit i8; exact f32 path,
3358            // bits still streamed once for both lanes.
3359            let (mut acc, mut nbits, mut idx) = (0u64, 0usize, 0usize);
3360            let (mut d1, mut d2) = (0f32, 0f32);
3361            for g in 0..ng {
3362                let so = (r * ng + g) * 2;
3363                let sgf = f16_to_f32(u16::from_le_bytes([
3364                    bytes[sc_off + so],
3365                    bytes[sc_off + so + 1],
3366                ]));
3367                let (mut g1, mut g2) = (0f32, 0f32);
3368                for k in 0..GROUP_SIZE {
3369                    if nbits < 8 {
3370                        acc = (acc << 8) | data[idx] as u64;
3371                        idx += 1;
3372                        nbits += 8;
3373                    }
3374                    let u = ((acc >> (nbits - 8)) & 0xFF) as i32;
3375                    nbits -= 8;
3376                    let w = (u - l) as f32;
3377                    g1 += w * x1[g * GROUP_SIZE + k];
3378                    g2 += w * x2[g * GROUP_SIZE + k];
3379                }
3380                d1 += g1 * sgf;
3381                d2 += g2 * sgf;
3382            }
3383            return (d1, d2);
3384        }
3385        thread_local! {
3386            static VBIT_SCRATCH2: std::cell::RefCell<Vec<u8>> =
3387                const { std::cell::RefCell::new(Vec::new()) };
3388        }
3389        #[inline(always)]
3390        fn fill<const B: usize>(data: &[u8], l: i32, buf: &mut [u8]) {
3391            for (blk, chunk) in buf.chunks_exact_mut(8).enumerate() {
3392                let u = unpack8::<B>(&data[blk * B..]);
3393                for k in 0..8 {
3394                    chunk[k] = (u[k] - l) as i8 as u8;
3395                }
3396            }
3397        }
3398        VBIT_SCRATCH2.with(|scratch| {
3399            let mut buf = scratch.borrow_mut();
3400            buf.resize(cols, 0);
3401            match b {
3402                3 => fill::<3>(data, l, &mut buf),
3403                4 => vbit_fill4(data, &mut buf),
3404                5 => fill::<5>(data, l, &mut buf),
3405                6 => fill::<6>(data, l, &mut buf),
3406                _ => unreachable!(),
3407            }
3408            let (mut d1, mut d2) = (0f32, 0f32);
3409            for g in 0..ng {
3410                let so = (r * ng + g) * 2;
3411                let s = f16_to_f32(u16::from_le_bytes([
3412                    bytes[sc_off + so],
3413                    bytes[sc_off + so + 1],
3414                ]));
3415                let wg = &buf[g * GROUP_SIZE..(g + 1) * GROUP_SIZE];
3416                let v1 = dot_i8_i8(wg, &a1.xq[g * GROUP_SIZE..(g + 1) * GROUP_SIZE]) as f32 * a1.sx;
3417                let v2 = dot_i8_i8(wg, &a2.xq[g * GROUP_SIZE..(g + 1) * GROUP_SIZE]) as f32 * a2.sx;
3418                d1 += v1 * s;
3419                d2 += v2 * s;
3420            }
3421            for &(j, xv) in &a1.outliers {
3422                let so = (r * ng + j / GROUP_SIZE) * 2;
3423                let s = f16_to_f32(u16::from_le_bytes([
3424                    bytes[sc_off + so],
3425                    bytes[sc_off + so + 1],
3426                ]));
3427                d1 += (buf[j] as i8) as f32 * s * xv;
3428            }
3429            for &(j, xv) in &a2.outliers {
3430                let so = (r * ng + j / GROUP_SIZE) * 2;
3431                let s = f16_to_f32(u16::from_le_bytes([
3432                    bytes[sc_off + so],
3433                    bytes[sc_off + so + 1],
3434                ]));
3435                d2 += (buf[j] as i8) as f32 * s * xv;
3436            }
3437            (d1, d2)
3438        })
3439    };
3440    for r in start..end {
3441        let (v1, v2) = row_dots(r);
3442        // SAFETY: disjoint row ranges per worker.
3443        unsafe {
3444            *p1.at(r) = v1;
3445            *p2.at(r) = v2;
3446        }
3447    }
3448}
3449
3450/// Two-input exact scalar vbit row range (same extraction) —
3451/// per-bit-width specialized, two accumulators per row; per-lane
3452/// accumulation order matches `vbitmatvec` exactly.
3453#[allow(clippy::too_many_arguments)]
3454fn vbit_range2_f32(
3455    bytes: &[u8],
3456    offsets: &[usize],
3457    x1: &[f32],
3458    x2: &[f32],
3459    rows: usize,
3460    cols: usize,
3461    p1: SendMut,
3462    p2: SendMut,
3463    start: usize,
3464    end: usize,
3465) {
3466    let ng = cols / GROUP_SIZE;
3467    let bits = &bytes[..rows];
3468    let sc_off = rows;
3469    #[inline(always)]
3470    #[allow(clippy::too_many_arguments)]
3471    fn dot_row2<const B: usize>(
3472        data: &[u8],
3473        bytes: &[u8],
3474        sc_off: usize,
3475        r: usize,
3476        ng: usize,
3477        x1: &[f32],
3478        x2: &[f32],
3479    ) -> (f32, f32) {
3480        let l = ((1i32 << (B - 1)) - 1) as f32;
3481        let gbytes = GROUP_SIZE * B / 8;
3482        let (mut d1, mut d2) = (0f32, 0f32);
3483        for g in 0..ng {
3484            let so = (r * ng + g) * 2;
3485            let s = f16_to_f32(u16::from_le_bytes([
3486                bytes[sc_off + so],
3487                bytes[sc_off + so + 1],
3488            ]));
3489            let x1g = &x1[g * GROUP_SIZE..(g + 1) * GROUP_SIZE];
3490            let x2g = &x2[g * GROUP_SIZE..(g + 1) * GROUP_SIZE];
3491            let gd0 = &data[g * gbytes..(g + 1) * gbytes];
3492            let (mut g1, mut g2) = (0f32, 0f32);
3493            for blk in 0..GROUP_SIZE / 8 {
3494                let u = unpack8::<B>(&gd0[blk * B..]);
3495                for k in 0..8 {
3496                    let w = u[k] as f32 - l;
3497                    g1 += w * x1g[blk * 8 + k];
3498                    g2 += w * x2g[blk * 8 + k];
3499                }
3500            }
3501            d1 += g1 * s;
3502            d2 += g2 * s;
3503        }
3504        (d1, d2)
3505    }
3506    for r in start..end {
3507        let data = &bytes[offsets[r]..offsets[r + 1]];
3508        let (v1, v2) = match bits[r] {
3509            3 => dot_row2::<3>(data, bytes, sc_off, r, ng, x1, x2),
3510            4 => dot_row2::<4>(data, bytes, sc_off, r, ng, x1, x2),
3511            5 => dot_row2::<5>(data, bytes, sc_off, r, ng, x1, x2),
3512            6 => dot_row2::<6>(data, bytes, sc_off, r, ng, x1, x2),
3513            8 => dot_row2::<8>(data, bytes, sc_off, r, ng, x1, x2),
3514            b => unreachable!("vbit bit-width {b} (validated at load)"),
3515        };
3516        // SAFETY: disjoint row ranges per worker.
3517        unsafe {
3518            *p1.at(r) = v1;
3519            *p2.at(r) = v2;
3520        }
3521    }
3522}
3523
3524// ───────────────────── q4_tiled kernels (§4.3) ─────────────────────
3525
3526/// One q4_tiled row dot on the A8W8 int8 path: per 32-group the tile
3527/// is ONE sequential read — [f16 scale][16B nibbles] — versus the two
3528/// distant streams of the split layout. Values/order identical to the
3529/// split kernels.
3530#[inline]
3531#[allow(unreachable_code)]
3532fn dot_q4t_row_i8(bytes: &[u8], r: usize, gpr: usize, xq: &[i8]) -> f32 {
3533    #[cfg(target_arch = "aarch64")]
3534    unsafe {
3535        return dot_q4t_row_sdot(bytes, r, gpr, xq);
3536    }
3537    #[cfg(target_arch = "x86_64")]
3538    unsafe {
3539        if vnni_tiles_enabled() {
3540            return dot_q4t_row_vnni(bytes, r, gpr, xq);
3541        }
3542        return dot_q4t_row_avx2(bytes, r, gpr, xq);
3543    }
3544    let mut acc = 0f32;
3545    for gi in 0..gpr {
3546        let tile = &bytes[(r * gpr + gi) * Q4_TILE..(r * gpr + gi + 1) * Q4_TILE];
3547        let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
3548        let mut d = 0i32;
3549        for (k, &b) in tile[2..].iter().enumerate() {
3550            d += ((b & 0x0F) as i32 - 8) * xq[gi * GROUP_SIZE + k * 2] as i32
3551                + (((b >> 4) & 0x0F) as i32 - 8) * xq[gi * GROUP_SIZE + k * 2 + 1] as i32;
3552        }
3553        acc += d as f32 * s;
3554    }
3555    acc
3556}
3557
3558#[cfg(target_arch = "aarch64")]
3559#[target_feature(enable = "neon,dotprod")]
3560unsafe fn dot_q4t_row_sdot(bytes: &[u8], r: usize, gpr: usize, xq: &[i8]) -> f32 {
3561    // SAFETY: callers uphold slice-length contracts (18B tile per group,
3562    // xq.len() == gpr·GROUP_SIZE).
3563    unsafe {
3564        use core::arch::aarch64::*;
3565        use core::arch::asm;
3566        let lomask = vdupq_n_u8(0x0F);
3567        let eight = vdupq_n_s8(8);
3568        let mut acc = 0f32;
3569        for gi in 0..gpr {
3570            let t = bytes.as_ptr().add((r * gpr + gi) * Q4_TILE);
3571            let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
3572            let b = vld1q_u8(t.add(2));
3573            let lo = vandq_u8(b, lomask);
3574            let hi = vshrq_n_u8::<4>(b);
3575            let e0 = vsubq_s8(vreinterpretq_s8_u8(vzip1q_u8(lo, hi)), eight);
3576            let e1 = vsubq_s8(vreinterpretq_s8_u8(vzip2q_u8(lo, hi)), eight);
3577            let x0 = vld1q_s8(xq.as_ptr().add(gi * GROUP_SIZE));
3578            let x1 = vld1q_s8(xq.as_ptr().add(gi * GROUP_SIZE + 16));
3579            let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
3580            asm!(
3581                "sdot {a0:v}.4s, {e0:v}.16b, {x0:v}.16b",
3582                "sdot {a1:v}.4s, {e1:v}.16b, {x1:v}.16b",
3583                a0 = inout(vreg) a0, a1 = inout(vreg) a1,
3584                e0 = in(vreg) e0, x0 = in(vreg) x0, e1 = in(vreg) e1, x1 = in(vreg) x1,
3585                options(pure, nomem, nostack),
3586            );
3587            acc += vaddvq_s32(vaddq_s32(a0, a1)) as f32 * s;
3588        }
3589        acc
3590    }
3591}
3592
3593#[cfg(target_arch = "x86_64")]
3594#[target_feature(enable = "avx2")]
3595unsafe fn dot_q4t_row_avx2(bytes: &[u8], r: usize, gpr: usize, xq: &[i8]) -> f32 {
3596    // SAFETY: see dot_q4t_row_sdot.
3597    unsafe {
3598        use core::arch::x86_64::*;
3599        let lomask = _mm_set1_epi8(0x0F);
3600        let eight = _mm256_set1_epi8(8);
3601        let ones = _mm256_set1_epi16(1);
3602        let mut acc = 0f32;
3603        for gi in 0..gpr {
3604            let t = bytes.as_ptr().add((r * gpr + gi) * Q4_TILE);
3605            let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
3606            let b = _mm_loadu_si128(t.add(2) as *const __m128i);
3607            let lo = _mm_and_si128(b, lomask);
3608            let hi = _mm_and_si128(_mm_srli_epi16::<4>(b), lomask);
3609            let w = _mm256_sub_epi8(
3610                _mm256_set_m128i(_mm_unpackhi_epi8(lo, hi), _mm_unpacklo_epi8(lo, hi)),
3611                eight,
3612            );
3613            let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
3614            let p16 = _mm256_maddubs_epi16(_mm256_abs_epi8(w), _mm256_sign_epi8(x, w));
3615            let d = _mm256_madd_epi16(p16, ones);
3616            let hi128 = _mm256_extracti128_si256::<1>(d);
3617            let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
3618            let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
3619            let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
3620            acc += _mm_cvtsi128_si32(s32) as f32 * s;
3621        }
3622        acc
3623    }
3624}
3625
3626/// VNNI twin of `dot_q4t_row_avx2`: same unpack, `vpdpbusd` replaces
3627/// the maddubs+madd pair (see `dpbusd_hsum` — sums are bit-identical).
3628/// 256-bit VL encoding, so the VEX `vpsignb` stays usable.
3629#[cfg(target_arch = "x86_64")]
3630#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
3631unsafe fn dot_q4t_row_vnni(bytes: &[u8], r: usize, gpr: usize, xq: &[i8]) -> f32 {
3632    // SAFETY: see dot_q4t_row_sdot.
3633    unsafe {
3634        use core::arch::x86_64::*;
3635        let lomask = _mm_set1_epi8(0x0F);
3636        let eight = _mm256_set1_epi8(8);
3637        let mut acc = 0f32;
3638        for gi in 0..gpr {
3639            let t = bytes.as_ptr().add((r * gpr + gi) * Q4_TILE);
3640            let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
3641            let b = _mm_loadu_si128(t.add(2) as *const __m128i);
3642            let lo = _mm_and_si128(b, lomask);
3643            let hi = _mm_and_si128(_mm_srli_epi16::<4>(b), lomask);
3644            let w = _mm256_sub_epi8(
3645                _mm256_set_m128i(_mm_unpackhi_epi8(lo, hi), _mm_unpacklo_epi8(lo, hi)),
3646                eight,
3647            );
3648            let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
3649            let d = dpbusd_hsum(_mm256_abs_epi8(w), _mm256_sign_epi8(x, w));
3650            acc += d as f32 * s;
3651        }
3652        acc
3653    }
3654}
3655
3656/// One q4_tiled row against FOUR activation streams: the nibble unpack
3657/// and abs() happen once per group instead of once per (group,
3658/// activation) — the unpack is the dominant per-element cost of the
3659/// tiled format (roadmap P0 portable blocking, q4t leg).
3660#[cfg(target_arch = "x86_64")]
3661// `fma` is NOT implied by `avx2`: without it LLVM lowers _mm256_fmadd_ps
3662// to a libm call per lane — measured 2x slower than the reduction this
3663// kernel replaces. The runtime gate (`avx2_enabled`) already requires
3664// both features, so declaring it here is safe.
3665#[target_feature(enable = "avx2,fma")]
3666unsafe fn dot_q4t_row_1x4_avx2(bytes: &[u8], r: usize, gpr: usize, xs: [&[i8]; 4]) -> [f32; 4] {
3667    // SAFETY: callers uphold the 18B-tile and xq-length contracts.
3668    unsafe {
3669        use core::arch::x86_64::*;
3670        let lomask = _mm_set1_epi8(0x0F);
3671        let eight = _mm256_set1_epi8(8);
3672        let ones = _mm256_set1_epi16(1);
3673        // One f32 accumulator VECTOR per activation, reduced once at the
3674        // end. Folding each group's i32 lanes to a scalar inside the loop
3675        // costs an extracti128 + three shift/add + a movd — a cross-lane
3676        // dependency chain per (group, activation), 288 of them per row at
3677        // cols=2304. The per-group scale is what forces a float
3678        // accumulator; it does not force a horizontal sum.
3679        //
3680        // The four accumulators are NAMED, not an array: as `[__m256; 4]`
3681        // indexed by a loop variable LLVM keeps them in memory and every
3682        // group pays four 32-byte loads and stores. That alone made this
3683        // kernel 2x SLOWER than the per-group reduction it replaces
3684        // (measured on the EPYC box: 150 s vs 71 s for two 256² steps).
3685        let mut f0 = _mm256_setzero_ps();
3686        let mut f1 = _mm256_setzero_ps();
3687        let mut f2 = _mm256_setzero_ps();
3688        let mut f3 = _mm256_setzero_ps();
3689        for gi in 0..gpr {
3690            let t = bytes.as_ptr().add((r * gpr + gi) * Q4_TILE);
3691            let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
3692            let sv = _mm256_set1_ps(s);
3693            let bb = _mm_loadu_si128(t.add(2) as *const __m128i);
3694            let lo = _mm_and_si128(bb, lomask);
3695            let hi = _mm_and_si128(_mm_srli_epi16::<4>(bb), lomask);
3696            let w = _mm256_sub_epi8(
3697                _mm256_set_m128i(_mm_unpackhi_epi8(lo, hi), _mm_unpacklo_epi8(lo, hi)),
3698                eight,
3699            );
3700            let aw = _mm256_abs_epi8(w);
3701            let off = gi * GROUP_SIZE;
3702            let dot = |xq: &[i8]| {
3703                let x = _mm256_loadu_si256(xq.as_ptr().add(off) as *const __m256i);
3704                let p16 = _mm256_maddubs_epi16(aw, _mm256_sign_epi8(x, w));
3705                _mm256_cvtepi32_ps(_mm256_madd_epi16(p16, ones))
3706            };
3707            f0 = _mm256_fmadd_ps(dot(xs[0]), sv, f0);
3708            f1 = _mm256_fmadd_ps(dot(xs[1]), sv, f1);
3709            f2 = _mm256_fmadd_ps(dot(xs[2]), sv, f2);
3710            f3 = _mm256_fmadd_ps(dot(xs[3]), sv, f3);
3711        }
3712        [
3713            hsum256_ps(f0),
3714            hsum256_ps(f1),
3715            hsum256_ps(f2),
3716            hsum256_ps(f3),
3717        ]
3718    }
3719}
3720
3721/// Horizontal sum of eight f32 lanes — the one cross-lane reduction the
3722/// blocked kernels pay, once per row instead of once per group.
3723#[cfg(target_arch = "x86_64")]
3724#[target_feature(enable = "avx2")]
3725#[inline]
3726unsafe fn hsum256_ps(v: core::arch::x86_64::__m256) -> f32 {
3727    // SAFETY: pure register arithmetic on the caller's vector.
3728    unsafe {
3729        use core::arch::x86_64::*;
3730        let hi = _mm256_extractf128_ps::<1>(v);
3731        let s = _mm_add_ps(_mm256_castps256_ps128(v), hi);
3732        let s = _mm_add_ps(s, _mm_movehl_ps(s, s));
3733        let s = _mm_add_ss(s, _mm_shuffle_ps::<0x55>(s, s));
3734        _mm_cvtss_f32(s)
3735    }
3736}
3737
3738/// VNNI twin of `dot_q4t_row_1x4_avx2` (see `dpbusd_hsum`).
3739#[cfg(target_arch = "x86_64")]
3740#[target_feature(enable = "avx2,fma,avx512f,avx512bw,avx512vl,avx512vnni")]
3741unsafe fn dot_q4t_row_1x4_vnni(bytes: &[u8], r: usize, gpr: usize, xs: [&[i8]; 4]) -> [f32; 4] {
3742    // SAFETY: callers uphold the 18B-tile and xq-length contracts.
3743    unsafe {
3744        use core::arch::x86_64::*;
3745        let lomask = _mm_set1_epi8(0x0F);
3746        let eight = _mm256_set1_epi8(8);
3747        // Same shape as the AVX2 twin: accumulate in f32 vectors and pay
3748        // one cross-lane reduction per row, not per (group, activation).
3749        let mut f0 = _mm256_setzero_ps();
3750        let mut f1 = _mm256_setzero_ps();
3751        let mut f2 = _mm256_setzero_ps();
3752        let mut f3 = _mm256_setzero_ps();
3753        for gi in 0..gpr {
3754            let t = bytes.as_ptr().add((r * gpr + gi) * Q4_TILE);
3755            let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
3756            let sv = _mm256_set1_ps(s);
3757            let bb = _mm_loadu_si128(t.add(2) as *const __m128i);
3758            let lo = _mm_and_si128(bb, lomask);
3759            let hi = _mm_and_si128(_mm_srli_epi16::<4>(bb), lomask);
3760            let w = _mm256_sub_epi8(
3761                _mm256_set_m128i(_mm_unpackhi_epi8(lo, hi), _mm_unpacklo_epi8(lo, hi)),
3762                eight,
3763            );
3764            let aw = _mm256_abs_epi8(w);
3765            let off = gi * GROUP_SIZE;
3766            let dot = |xq: &[i8]| {
3767                let x = _mm256_loadu_si256(xq.as_ptr().add(off) as *const __m256i);
3768                _mm256_cvtepi32_ps(_mm256_dpbusd_epi32(
3769                    _mm256_setzero_si256(),
3770                    aw,
3771                    _mm256_sign_epi8(x, w),
3772                ))
3773            };
3774            f0 = _mm256_fmadd_ps(dot(xs[0]), sv, f0);
3775            f1 = _mm256_fmadd_ps(dot(xs[1]), sv, f1);
3776            f2 = _mm256_fmadd_ps(dot(xs[2]), sv, f2);
3777            f3 = _mm256_fmadd_ps(dot(xs[3]), sv, f3);
3778        }
3779        let acc = [
3780            hsum256_ps(f0),
3781            hsum256_ps(f1),
3782            hsum256_ps(f2),
3783            hsum256_ps(f3),
3784        ];
3785        acc
3786    }
3787}
3788
3789/// ARM twin of `dot_q4t_row_1x4_avx2`: one nibble unpack per group
3790/// serves FOUR activation streams. Per stream the group order and f32
3791/// accumulation match `dot_q4t_row_sdot` exactly — batch == matvec
3792/// bit-for-bit.
3793#[cfg(target_arch = "aarch64")]
3794#[target_feature(enable = "neon,dotprod")]
3795unsafe fn dot_q4t_row_1x4_sdot(bytes: &[u8], r: usize, gpr: usize, xs: [&[i8]; 4]) -> [f32; 4] {
3796    // SAFETY: callers uphold the 18B-tile and xq-length contracts.
3797    unsafe {
3798        use core::arch::aarch64::*;
3799        use core::arch::asm;
3800        let lomask = vdupq_n_u8(0x0F);
3801        let eight = vdupq_n_s8(8);
3802        let mut acc = [0f32; 4];
3803        for gi in 0..gpr {
3804            let t = bytes.as_ptr().add((r * gpr + gi) * Q4_TILE);
3805            let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
3806            let b = vld1q_u8(t.add(2));
3807            let lo = vandq_u8(b, lomask);
3808            let hi = vshrq_n_u8::<4>(b);
3809            let e0 = vsubq_s8(vreinterpretq_s8_u8(vzip1q_u8(lo, hi)), eight);
3810            let e1 = vsubq_s8(vreinterpretq_s8_u8(vzip2q_u8(lo, hi)), eight);
3811            for (k, xq) in xs.iter().enumerate() {
3812                let x0 = vld1q_s8(xq.as_ptr().add(gi * GROUP_SIZE));
3813                let x1 = vld1q_s8(xq.as_ptr().add(gi * GROUP_SIZE + 16));
3814                let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
3815                asm!(
3816                    "sdot {a0:v}.4s, {e0:v}.16b, {x0:v}.16b",
3817                    "sdot {a1:v}.4s, {e1:v}.16b, {x1:v}.16b",
3818                    a0 = inout(vreg) a0, a1 = inout(vreg) a1,
3819                    e0 = in(vreg) e0, x0 = in(vreg) x0, e1 = in(vreg) e1, x1 = in(vreg) x1,
3820                    options(pure, nomem, nostack),
3821                );
3822                acc[k] += vaddvq_s32(vaddq_s32(a0, a1)) as f32 * s;
3823            }
3824        }
3825        acc
3826    }
3827}
3828
3829/// Exact-term correction for A8W8 outliers on a tiled row.
3830#[inline]
3831fn q4t_outlier(bytes: &[u8], r: usize, gpr: usize, j: usize) -> (f32, f32) {
3832    let gi = j / GROUP_SIZE;
3833    let k = j % GROUP_SIZE;
3834    let tile = &bytes[(r * gpr + gi) * Q4_TILE..(r * gpr + gi + 1) * Q4_TILE];
3835    let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
3836    let byte = tile[2 + k / 2];
3837    let nib = if k & 1 == 0 { byte & 0x0F } else { byte >> 4 };
3838    ((nib as i32 - 8) as f32, s)
3839}
3840
3841/// Exact scalar q4_tiled row (CMF_SDOT=0 contract) — same pairwise
3842/// accumulation shape as `q4_range_f32`.
3843#[inline]
3844fn q4t_row_exact(bytes: &[u8], r: usize, gpr: usize, x: &[f32]) -> f32 {
3845    let mut acc = 0f32;
3846    for gi in 0..gpr {
3847        let tile = &bytes[(r * gpr + gi) * Q4_TILE..(r * gpr + gi + 1) * Q4_TILE];
3848        let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
3849        let xg = &x[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
3850        let mut ga = 0f32;
3851        for (k, &b) in tile[2..].iter().enumerate() {
3852            ga += ((b & 0x0F) as f32 - 8.0) * xg[k * 2]
3853                + (((b >> 4) & 0x0F) as f32 - 8.0) * xg[k * 2 + 1];
3854        }
3855        acc += ga * s;
3856    }
3857    acc
3858}
3859
3860/// Split view of a `q4tp` payload. The three planes are resolved once per
3861/// matvec instead of per row — `q4tp_sections` is cheap, but doing it inside
3862/// the row loop would put a division on the hot path for nothing.
3863struct Q4tpView<'a> {
3864    nib: &'a [u8],
3865    params: &'a [u8],
3866    codes: &'a [u8],
3867    stride: usize,
3868    /// q2tp reads the ladder with rung 0 = exact zero.
3869    zero_rung: bool,
3870}
3871
3872impl<'a> Q4tpView<'a> {
3873    fn new(bytes: &'a [u8], rows: usize, cols: usize) -> Self {
3874        let (params_off, codes_off, stride) = q4tp_sections(rows, cols);
3875        Self {
3876            nib: &bytes[..params_off],
3877            params: &bytes[params_off..codes_off],
3878            codes: &bytes[codes_off..],
3879            stride,
3880            zero_rung: false,
3881        }
3882    }
3883
3884    /// The q2tp view: identical params/codes planes, 8 B weight chunks.
3885    fn new_q2(bytes: &'a [u8], rows: usize, cols: usize) -> Self {
3886        let (params_off, codes_off, stride) = q2tp_sections(rows, cols);
3887        Self {
3888            nib: &bytes[..params_off],
3889            params: &bytes[params_off..codes_off],
3890            codes: &bytes[codes_off..],
3891            stride,
3892            zero_rung: true,
3893        }
3894    }
3895
3896    /// Expand row `r`'s per-tile scales into `out` (length `gpr`).
3897    ///
3898    /// Doing this once per row — rather than decoding a 5-bit code inside the
3899    /// tile loop — is what makes the format free at runtime. Random access to
3900    /// a packed 5-bit field costs a division, two bounds checks and a branch;
3901    /// the tile's actual work is two `sdot`s, so per-tile decoding dominated
3902    /// the kernel and cost 5x (measured: 1.4 vs 6.9 tok/s on Nanbeige-3B).
3903    /// Walking the plane sequentially with a bit accumulator is ~3 ops.
3904    /// Eight 5-bit codes are exactly five bytes, so a whole group of
3905    /// eight decodes from one little-endian word at fixed shifts. The
3906    /// bit-accumulator this replaces carried a data-dependent `while
3907    /// have < 5` refill whose branch sat in the innermost loop of every
3908    /// q4tp row; a decode profile put this function above the dot
3909    /// products it feeds. Same bitstream, same codes — just no branch
3910    /// and eight independent extractions.
3911    #[inline]
3912    fn scales_into(&self, r: usize, gpr: usize, out: &mut [f32]) {
3913        let tab = if self.zero_rung {
3914            q2tp_ladder(self.params, r)
3915        } else {
3916            q4tp_ladder(self.params, r)
3917        };
3918        let codes = &self.codes[r * self.stride..(r + 1) * self.stride];
3919        let out = &mut out[..gpr];
3920        let mut chunks = out.chunks_exact_mut(8);
3921        let mut ci = 0usize;
3922        for c in &mut chunks {
3923            let w = u64::from(codes[ci])
3924                | u64::from(codes[ci + 1]) << 8
3925                | u64::from(codes[ci + 2]) << 16
3926                | u64::from(codes[ci + 3]) << 24
3927                | u64::from(codes[ci + 4]) << 32;
3928            for (k, o) in c.iter_mut().enumerate() {
3929                *o = tab[((w >> (5 * k)) & 31) as usize];
3930            }
3931            ci += 5;
3932        }
3933        // Fewer than eight codes left: the shared total accessor, which
3934        // tolerates a 5-bit field whose spill byte is past the stride.
3935        let tail = &codes[ci..];
3936        for (k, o) in chunks.into_remainder().iter_mut().enumerate() {
3937            *o = tab[q4tp_code(tail, k)];
3938        }
3939    }
3940}
3941
3942#[inline]
3943fn dot_q4tp_row_i8(nib: &[u8], r: usize, gpr: usize, xq: &[i8], scales: &[f32]) -> f32 {
3944    #[cfg(target_arch = "aarch64")]
3945    unsafe {
3946        return dot_q4tp_row_sdot(nib, r, gpr, xq, scales);
3947    }
3948    #[cfg(target_arch = "x86_64")]
3949    unsafe {
3950        if vnni_tiles_enabled() {
3951            return dot_q4tp_row_vnni(nib, r, gpr, xq, scales);
3952        }
3953        return dot_q4tp_row_avx2(nib, r, gpr, xq, scales);
3954    }
3955    #[allow(unreachable_code)]
3956    {
3957        let mut acc = 0f32;
3958        for gi in 0..gpr {
3959            let tile = &nib[(r * gpr + gi) * Q4TP_NIB..(r * gpr + gi + 1) * Q4TP_NIB];
3960            let s = scales[gi];
3961            let mut d = 0i32;
3962            for (k, &b) in tile.iter().enumerate() {
3963                d += ((b & 0x0F) as i32 - 8) * xq[gi * GROUP_SIZE + k * 2] as i32
3964                    + (((b >> 4) & 0x0F) as i32 - 8) * xq[gi * GROUP_SIZE + k * 2 + 1] as i32;
3965            }
3966            acc += d as f32 * s;
3967        }
3968        acc
3969    }
3970}
3971
3972/// q4tp twin of `dot_q4t_row_sdot`: identical nibble math, but the tile
3973/// stride is 16 B (no inline scale) and the scale is a ladder lookup.
3974#[cfg(target_arch = "aarch64")]
3975#[target_feature(enable = "neon,dotprod")]
3976unsafe fn dot_q4tp_row_sdot(nib: &[u8], r: usize, gpr: usize, xq: &[i8], scales: &[f32]) -> f32 {
3977    // SAFETY: callers uphold slice-length contracts (16B tile per group,
3978    // xq.len() == gpr·GROUP_SIZE, codes covering gpr 5-bit fields).
3979    unsafe {
3980        use core::arch::aarch64::*;
3981        use core::arch::asm;
3982        let lomask = vdupq_n_u8(0x0F);
3983        let eight = vdupq_n_s8(8);
3984        let mut acc = 0f32;
3985        for gi in 0..gpr {
3986            let t = nib.as_ptr().add((r * gpr + gi) * Q4TP_NIB);
3987            let s = *scales.get_unchecked(gi);
3988            let b = vld1q_u8(t);
3989            let lo = vandq_u8(b, lomask);
3990            let hi = vshrq_n_u8::<4>(b);
3991            let e0 = vsubq_s8(vreinterpretq_s8_u8(vzip1q_u8(lo, hi)), eight);
3992            let e1 = vsubq_s8(vreinterpretq_s8_u8(vzip2q_u8(lo, hi)), eight);
3993            let x0 = vld1q_s8(xq.as_ptr().add(gi * GROUP_SIZE));
3994            let x1 = vld1q_s8(xq.as_ptr().add(gi * GROUP_SIZE + 16));
3995            let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
3996            asm!(
3997                "sdot {a0:v}.4s, {e0:v}.16b, {x0:v}.16b",
3998                "sdot {a1:v}.4s, {e1:v}.16b, {x1:v}.16b",
3999                a0 = inout(vreg) a0, a1 = inout(vreg) a1,
4000                e0 = in(vreg) e0, x0 = in(vreg) x0, e1 = in(vreg) e1, x1 = in(vreg) x1,
4001                options(pure, nomem, nostack),
4002            );
4003            acc += vaddvq_s32(vaddq_s32(a0, a1)) as f32 * s;
4004        }
4005        acc
4006    }
4007}
4008
4009#[cfg(target_arch = "x86_64")]
4010#[target_feature(enable = "avx2")]
4011unsafe fn dot_q4tp_row_avx2(nib: &[u8], r: usize, gpr: usize, xq: &[i8], scales: &[f32]) -> f32 {
4012    // SAFETY: see dot_q4tp_row_sdot.
4013    unsafe {
4014        use core::arch::x86_64::*;
4015        let lomask = _mm_set1_epi8(0x0F);
4016        let eight = _mm256_set1_epi8(8);
4017        let ones = _mm256_set1_epi16(1);
4018        let mut acc = 0f32;
4019        for gi in 0..gpr {
4020            let t = nib.as_ptr().add((r * gpr + gi) * Q4TP_NIB);
4021            let s = *scales.get_unchecked(gi);
4022            let b = _mm_loadu_si128(t as *const __m128i);
4023            let lo = _mm_and_si128(b, lomask);
4024            let hi = _mm_and_si128(_mm_srli_epi16::<4>(b), lomask);
4025            let w = _mm256_sub_epi8(
4026                _mm256_set_m128i(_mm_unpackhi_epi8(lo, hi), _mm_unpacklo_epi8(lo, hi)),
4027                eight,
4028            );
4029            let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
4030            let p16 = _mm256_maddubs_epi16(_mm256_abs_epi8(w), _mm256_sign_epi8(x, w));
4031            let d = _mm256_madd_epi16(p16, ones);
4032            let hi128 = _mm256_extracti128_si256::<1>(d);
4033            let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
4034            let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
4035            let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
4036            acc += _mm_cvtsi128_si32(s32) as f32 * s;
4037        }
4038        acc
4039    }
4040}
4041
4042/// VNNI twin of `dot_q4tp_row_avx2` (see `dot_q4t_row_vnni` for why the
4043/// 256-bit VL encoding is the one to use here).
4044#[cfg(target_arch = "x86_64")]
4045#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
4046unsafe fn dot_q4tp_row_vnni(nib: &[u8], r: usize, gpr: usize, xq: &[i8], scales: &[f32]) -> f32 {
4047    // SAFETY: see dot_q4tp_row_sdot.
4048    unsafe {
4049        use core::arch::x86_64::*;
4050        let lomask = _mm_set1_epi8(0x0F);
4051        let eight = _mm256_set1_epi8(8);
4052        let mut acc = 0f32;
4053        for gi in 0..gpr {
4054            let t = nib.as_ptr().add((r * gpr + gi) * Q4TP_NIB);
4055            let s = *scales.get_unchecked(gi);
4056            let b = _mm_loadu_si128(t as *const __m128i);
4057            let lo = _mm_and_si128(b, lomask);
4058            let hi = _mm_and_si128(_mm_srli_epi16::<4>(b), lomask);
4059            let w = _mm256_sub_epi8(
4060                _mm256_set_m128i(_mm_unpackhi_epi8(lo, hi), _mm_unpacklo_epi8(lo, hi)),
4061                eight,
4062            );
4063            let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
4064            acc += dpbusd_hsum(_mm256_abs_epi8(w), _mm256_sign_epi8(x, w)) as f32 * s;
4065        }
4066        acc
4067    }
4068}
4069
4070/// Exact scalar q4tp row — the `CMF_SDOT=0` contract, same pairwise
4071/// accumulation shape as `q4t_row_exact`.
4072#[inline]
4073fn q4tp_row_exact(nib: &[u8], r: usize, gpr: usize, x: &[f32], scales: &[f32]) -> f32 {
4074    let mut acc = 0f32;
4075    for gi in 0..gpr {
4076        let tile = &nib[(r * gpr + gi) * Q4TP_NIB..(r * gpr + gi + 1) * Q4TP_NIB];
4077        let s = scales[gi];
4078        let xg = &x[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
4079        let mut ga = 0f32;
4080        for (k, &b) in tile.iter().enumerate() {
4081            ga += ((b & 0x0F) as f32 - 8.0) * xg[k * 2]
4082                + (((b >> 4) & 0x0F) as f32 - 8.0) * xg[k * 2 + 1];
4083        }
4084        acc += ga * s;
4085    }
4086    acc
4087}
4088
4089/// Single weight of a q4tp tensor — the a8w8 outlier path, which restores
4090/// activation outliers at full precision after the int8 pass.
4091#[inline]
4092fn q4tp_outlier(nib: &[u8], r: usize, gpr: usize, j: usize, scales: &[f32]) -> (f32, f32) {
4093    let (gi, k) = (j / GROUP_SIZE, j % GROUP_SIZE);
4094    let byte = nib[(r * gpr + gi) * Q4TP_NIB + k / 2];
4095    let n = if k & 1 == 0 { byte & 0x0F } else { byte >> 4 };
4096    ((n as i32 - 8) as f32, scales[gi])
4097}
4098
4099/// Fused q4tp matvec (dispatch mirrors `q4t_matvec`).
4100fn q4tp_matvec(
4101    bytes: &[u8],
4102    x: &[f32],
4103    rows: usize,
4104    cols: usize,
4105    out: &mut [f32],
4106    pool: Option<&Pool>,
4107) {
4108    debug_assert_eq!(out.len(), rows);
4109    let gpr = cols / GROUP_SIZE;
4110    let v = Q4tpView::new(bytes, rows, cols);
4111    let out_addr = SendMut(out.as_mut_ptr());
4112    if a8w8_enabled() {
4113        let act = split_act(x);
4114        let run = |start: usize, end: usize| {
4115            // One scratch row of scales per worker, reused across its rows.
4116            let mut sc = vec![0f32; gpr];
4117            for r in start..end {
4118                v.scales_into(r, gpr, &mut sc);
4119                let mut acc = dot_q4tp_row_i8(v.nib, r, gpr, &act.xq, &sc) * act.sx;
4120                for &(j, xv) in &act.outliers {
4121                    let (w, s) = q4tp_outlier(v.nib, r, gpr, j, &sc);
4122                    acc += w * s * xv;
4123                }
4124                // SAFETY: disjoint row ranges per worker.
4125                unsafe { *out_addr.at(r) = acc };
4126            }
4127        };
4128        dispatch_rows(pool, rows, &run);
4129        return;
4130    }
4131    let run = |start: usize, end: usize| {
4132        let mut sc = vec![0f32; gpr];
4133        for r in start..end {
4134            v.scales_into(r, gpr, &mut sc);
4135            // SAFETY: disjoint row ranges per worker.
4136            unsafe { *out_addr.at(r) = q4tp_row_exact(v.nib, r, gpr, x, &sc) };
4137        }
4138    };
4139    dispatch_rows(pool, rows, &run);
4140}
4141
4142/// Fused two-input q4tp matvec — the SwiGLU gate/up pair. Weights and the
4143/// row ladder are read once and spent on both activation streams.
4144#[allow(clippy::too_many_arguments)]
4145fn q4tp_matvec2(
4146    bytes: &[u8],
4147    x1: &[f32],
4148    x2: &[f32],
4149    rows: usize,
4150    cols: usize,
4151    o1: &mut [f32],
4152    o2: &mut [f32],
4153    pool: Option<&Pool>,
4154) {
4155    let gpr = cols / GROUP_SIZE;
4156    let v = Q4tpView::new(bytes, rows, cols);
4157    let (p1, p2) = (SendMut(o1.as_mut_ptr()), SendMut(o2.as_mut_ptr()));
4158    let run = |start: usize, end: usize| {
4159        let mut sc = vec![0f32; gpr];
4160        for r in start..end {
4161            v.scales_into(r, gpr, &mut sc);
4162            // SAFETY: disjoint row ranges per worker.
4163            unsafe {
4164                *p1.at(r) = q4tp_row_exact(v.nib, r, gpr, x1, &sc);
4165                *p2.at(r) = q4tp_row_exact(v.nib, r, gpr, x2, &sc);
4166            }
4167        }
4168    };
4169    dispatch_rows(pool, rows, &run);
4170}
4171
4172/// Exact f32 dot of one q2tp row: 2-bit fields LSB-first, (c − 1.5)·s.
4173/// Scalar on purpose — the 2-bit class targets the GPU graph; the CPU
4174/// path exists for parity gates and small-machine fallback.
4175fn q2tp_row_exact(chunks: &[u8], r: usize, gpr: usize, x: &[f32], scales: &[f32]) -> f32 {
4176    let mut acc = 0f32;
4177    for gi in 0..gpr {
4178        let ch = &chunks[(r * gpr + gi) * Q2TP_CHUNK..(r * gpr + gi + 1) * Q2TP_CHUNK];
4179        let s = scales[gi];
4180        let xb = &x[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
4181        let mut g = 0f32;
4182        for (k, &b) in ch.iter().enumerate() {
4183            g += ((b & 3) as f32 - 1.5) * xb[k * 4]
4184                + (((b >> 2) & 3) as f32 - 1.5) * xb[k * 4 + 1]
4185                + (((b >> 4) & 3) as f32 - 1.5) * xb[k * 4 + 2]
4186                + (((b >> 6) & 3) as f32 - 1.5) * xb[k * 4 + 3];
4187        }
4188        acc += s * g;
4189    }
4190    acc
4191}
4192
4193fn q2tp_matvec(
4194    bytes: &[u8],
4195    x: &[f32],
4196    rows: usize,
4197    cols: usize,
4198    out: &mut [f32],
4199    pool: Option<&Pool>,
4200) {
4201    debug_assert_eq!(out.len(), rows);
4202    let gpr = cols / GROUP_SIZE;
4203    let v = Q4tpView::new_q2(bytes, rows, cols);
4204    let out_addr = SendMut(out.as_mut_ptr());
4205    let run = |start: usize, end: usize| {
4206        let mut sc = vec![0f32; gpr];
4207        for r in start..end {
4208            v.scales_into(r, gpr, &mut sc);
4209            // SAFETY: disjoint row ranges per worker.
4210            unsafe { *out_addr.at(r) = q2tp_row_exact(v.nib, r, gpr, x, &sc) };
4211        }
4212    };
4213    dispatch_rows(pool, rows, &run);
4214}
4215
4216/// Fused two-input q2tp matvec — the SwiGLU gate/up pair.
4217#[allow(clippy::too_many_arguments)]
4218fn q2tp_matvec2(
4219    bytes: &[u8],
4220    x1: &[f32],
4221    x2: &[f32],
4222    rows: usize,
4223    cols: usize,
4224    o1: &mut [f32],
4225    o2: &mut [f32],
4226    pool: Option<&Pool>,
4227) {
4228    let gpr = cols / GROUP_SIZE;
4229    let v = Q4tpView::new_q2(bytes, rows, cols);
4230    let (p1, p2) = (SendMut(o1.as_mut_ptr()), SendMut(o2.as_mut_ptr()));
4231    let run = |start: usize, end: usize| {
4232        let mut sc = vec![0f32; gpr];
4233        for r in start..end {
4234            v.scales_into(r, gpr, &mut sc);
4235            // SAFETY: disjoint row ranges per worker.
4236            unsafe {
4237                *p1.at(r) = q2tp_row_exact(v.nib, r, gpr, x1, &sc);
4238                *p2.at(r) = q2tp_row_exact(v.nib, r, gpr, x2, &sc);
4239            }
4240        }
4241    };
4242    dispatch_rows(pool, rows, &run);
4243}
4244
4245/// Batched q2tp matmat: scalar row kernel over every batch column. CPU
4246/// prefill only — decode rides the graph, so plain and correct beats
4247/// clever here.
4248fn q2tp_matmat(
4249    bytes: &[u8],
4250    xs_all: &[f32],
4251    b: usize,
4252    rows: usize,
4253    cols: usize,
4254    out: &mut [f32],
4255    pool: Option<&Pool>,
4256) {
4257    debug_assert_eq!(out.len(), b * rows);
4258    let gpr = cols / GROUP_SIZE;
4259    let v = Q4tpView::new_q2(bytes, rows, cols);
4260    let out_addr = SendMut(out.as_mut_ptr());
4261    let run = |start: usize, end: usize| {
4262        let mut sc = vec![0f32; gpr];
4263        for r in start..end {
4264            v.scales_into(r, gpr, &mut sc);
4265            for bi in 0..b {
4266                let x = &xs_all[bi * cols..(bi + 1) * cols];
4267                // SAFETY: disjoint row ranges per worker.
4268                unsafe { *out_addr.at(bi * rows + r) = q2tp_row_exact(v.nib, r, gpr, x, &sc) };
4269            }
4270        }
4271    };
4272    dispatch_rows(pool, rows, &run);
4273}
4274
4275/// Four batch columns against one q4tp row: the tile is unpacked ONCE and
4276/// spent on four activation streams, which is where a prefill batch stops
4277/// being weight-bandwidth-bound. Twin of `dot_q4t_row_1x4_sdot`.
4278#[cfg(target_arch = "aarch64")]
4279#[target_feature(enable = "neon,dotprod")]
4280unsafe fn dot_q4tp_row_1x4_sdot(
4281    nib: &[u8],
4282    r: usize,
4283    gpr: usize,
4284    xs: [&[i8]; 4],
4285    scales: &[f32],
4286) -> [f32; 4] {
4287    // SAFETY: see dot_q4tp_row_sdot; every xs[k] is gpr·GROUP_SIZE long.
4288    unsafe {
4289        use core::arch::aarch64::*;
4290        use core::arch::asm;
4291        let lomask = vdupq_n_u8(0x0F);
4292        let eight = vdupq_n_s8(8);
4293        // Named accumulators, NOT an array indexed by a loop variable: the
4294        // latter does not stay in registers (the same defect cost 2x in the
4295        // AVX2 q4t kernel and again in WGSL).
4296        let (mut f0, mut f1, mut f2, mut f3) = (0f32, 0f32, 0f32, 0f32);
4297        for gi in 0..gpr {
4298            let t = nib.as_ptr().add((r * gpr + gi) * Q4TP_NIB);
4299            let s = *scales.get_unchecked(gi);
4300            let bb = vld1q_u8(t);
4301            let lo = vandq_u8(bb, lomask);
4302            let hi = vshrq_n_u8::<4>(bb);
4303            let e0 = vsubq_s8(vreinterpretq_s8_u8(vzip1q_u8(lo, hi)), eight);
4304            let e1 = vsubq_s8(vreinterpretq_s8_u8(vzip2q_u8(lo, hi)), eight);
4305            let mut d = [0f32; 4];
4306            for (k, dk) in d.iter_mut().enumerate() {
4307                let x0 = vld1q_s8(xs[k].as_ptr().add(gi * GROUP_SIZE));
4308                let x1 = vld1q_s8(xs[k].as_ptr().add(gi * GROUP_SIZE + 16));
4309                let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
4310                asm!(
4311                    "sdot {a0:v}.4s, {e0:v}.16b, {x0:v}.16b",
4312                    "sdot {a1:v}.4s, {e1:v}.16b, {x1:v}.16b",
4313                    a0 = inout(vreg) a0, a1 = inout(vreg) a1,
4314                    e0 = in(vreg) e0, x0 = in(vreg) x0, e1 = in(vreg) e1, x1 = in(vreg) x1,
4315                    options(pure, nomem, nostack),
4316                );
4317                *dk = vaddvq_s32(vaddq_s32(a0, a1)) as f32 * s;
4318            }
4319            f0 += d[0];
4320            f1 += d[1];
4321            f2 += d[2];
4322            f3 += d[3];
4323        }
4324        [f0, f1, f2, f3]
4325    }
4326}
4327
4328/// Fused q4tp matmat — the same three arms `q4t_matmat` has. Shipping only
4329/// the scalar one made Nanbeige-3B decode at 1.2 tok/s against q4t's 5.9:
4330/// the format was fine, the missing arms were the whole regression.
4331fn q4tp_matmat(
4332    bytes: &[u8],
4333    xs_all: &[f32],
4334    b: usize,
4335    rows: usize,
4336    cols: usize,
4337    out: &mut [f32],
4338    pool: Option<&Pool>,
4339) {
4340    debug_assert_eq!(out.len(), b * rows);
4341    let gpr = cols / GROUP_SIZE;
4342    let v = Q4tpView::new(bytes, rows, cols);
4343
4344    // Wide batches ride the AMX through a dequant-tile sgemm, as in q4t.
4345    #[cfg(target_os = "macos")]
4346    if b >= 8 && rows * cols >= 500_000 && accel_gemm_enabled() {
4347        dequant_matmat_accel(
4348            &|r, dst| {
4349                let mut sc = [0f32; 32];
4350                let mut scv;
4351                let s: &[f32] = if gpr <= 32 {
4352                    v.scales_into(r, gpr, &mut sc);
4353                    &sc[..gpr]
4354                } else {
4355                    scv = vec![0f32; gpr];
4356                    v.scales_into(r, gpr, &mut scv);
4357                    &scv
4358                };
4359                for gi in 0..gpr {
4360                    let tile = &v.nib[(r * gpr + gi) * Q4TP_NIB..(r * gpr + gi + 1) * Q4TP_NIB];
4361                    for (k, &bb) in tile.iter().enumerate() {
4362                        dst[gi * GROUP_SIZE + k * 2] = ((bb & 0x0F) as f32 - 8.0) * s[gi];
4363                        dst[gi * GROUP_SIZE + k * 2 + 1] =
4364                            (((bb >> 4) & 0x0F) as f32 - 8.0) * s[gi];
4365                    }
4366                }
4367            },
4368            xs_all,
4369            b,
4370            rows,
4371            cols,
4372            out,
4373            pool,
4374        );
4375        return;
4376    }
4377
4378    let out_addr = SendMut(out.as_mut_ptr());
4379    if a8w8_enabled() {
4380        let acts: Vec<SplitAct> = (0..b)
4381            .map(|bi| split_act(&xs_all[bi * cols..(bi + 1) * cols]))
4382            .collect();
4383        let acts = &acts;
4384        #[cfg(target_arch = "aarch64")]
4385        let blocked_ok = sdot_enabled()
4386            && std::env::var("CMF_X86_BLOCKED")
4387                .map(|val| val != "0")
4388                .unwrap_or(true);
4389        #[cfg(not(target_arch = "aarch64"))]
4390        let blocked_ok = false;
4391        let run = |start: usize, end: usize| {
4392            let mut sc = vec![0f32; gpr];
4393            for r in start..end {
4394                v.scales_into(r, gpr, &mut sc);
4395                let mut bi = 0usize;
4396                #[cfg(target_arch = "aarch64")]
4397                if blocked_ok {
4398                    while bi + 4 <= acts.len() {
4399                        let xs = [
4400                            acts[bi].xq.as_slice(),
4401                            acts[bi + 1].xq.as_slice(),
4402                            acts[bi + 2].xq.as_slice(),
4403                            acts[bi + 3].xq.as_slice(),
4404                        ];
4405                        let d = unsafe { dot_q4tp_row_1x4_sdot(v.nib, r, gpr, xs, &sc) };
4406                        for k in 0..4 {
4407                            let act = &acts[bi + k];
4408                            let mut acc = d[k] * act.sx;
4409                            for &(j, xv) in &act.outliers {
4410                                let (w, s) = q4tp_outlier(v.nib, r, gpr, j, &sc);
4411                                acc += w * s * xv;
4412                            }
4413                            // SAFETY: disjoint (bi, r) cells per worker.
4414                            unsafe { *out_addr.at((bi + k) * rows + r) = acc };
4415                        }
4416                        bi += 4;
4417                    }
4418                }
4419                let _ = blocked_ok;
4420                while bi < acts.len() {
4421                    let act = &acts[bi];
4422                    let mut acc = dot_q4tp_row_i8(v.nib, r, gpr, &act.xq, &sc) * act.sx;
4423                    for &(j, xv) in &act.outliers {
4424                        let (w, s) = q4tp_outlier(v.nib, r, gpr, j, &sc);
4425                        acc += w * s * xv;
4426                    }
4427                    // SAFETY: disjoint (bi, r) cells per worker range.
4428                    unsafe { *out_addr.at(bi * rows + r) = acc };
4429                    bi += 1;
4430                }
4431            }
4432        };
4433        dispatch_rows(pool, rows, &run);
4434        return;
4435    }
4436
4437    let run = |start: usize, end: usize| {
4438        let mut sc = vec![0f32; gpr];
4439        for r in start..end {
4440            v.scales_into(r, gpr, &mut sc);
4441            for bi in 0..b {
4442                let x = &xs_all[bi * cols..(bi + 1) * cols];
4443                // SAFETY: disjoint (bi, r) cells per worker range.
4444                unsafe { *out_addr.at(bi * rows + r) = q4tp_row_exact(v.nib, r, gpr, x, &sc) };
4445            }
4446        }
4447    };
4448    dispatch_rows(pool, rows, &run);
4449}
4450
4451/// Fused q4_tiled matvec (dispatch mirrors `q4matvec`).
4452fn q4t_matvec(
4453    bytes: &[u8],
4454    x: &[f32],
4455    rows: usize,
4456    cols: usize,
4457    out: &mut [f32],
4458    pool: Option<&Pool>,
4459) {
4460    debug_assert_eq!(out.len(), rows);
4461    let gpr = cols / GROUP_SIZE;
4462    let out_addr = SendMut(out.as_mut_ptr());
4463    if a8w8_enabled() {
4464        let act = split_act(x);
4465        let run = move |start: usize, end: usize| {
4466            for r in start..end {
4467                let mut acc = dot_q4t_row_i8(bytes, r, gpr, &act.xq) * act.sx;
4468                for &(j, xv) in &act.outliers {
4469                    let (w, s) = q4t_outlier(bytes, r, gpr, j);
4470                    acc += w * s * xv;
4471                }
4472                // SAFETY: disjoint row ranges per worker.
4473                unsafe { *out_addr.at(r) = acc };
4474            }
4475        };
4476        dispatch_rows(pool, rows, &run);
4477        return;
4478    }
4479    let run = move |start: usize, end: usize| {
4480        for r in start..end {
4481            // SAFETY: disjoint row ranges per worker.
4482            unsafe { *out_addr.at(r) = q4t_row_exact(bytes, r, gpr, x) };
4483        }
4484    };
4485    dispatch_rows(pool, rows, &run);
4486}
4487
4488/// Fused two-input q4_tiled matvec (weights read once per pair).
4489#[allow(clippy::too_many_arguments)]
4490fn q4t_matvec2(
4491    bytes: &[u8],
4492    x1: &[f32],
4493    x2: &[f32],
4494    rows: usize,
4495    cols: usize,
4496    o1: &mut [f32],
4497    o2: &mut [f32],
4498    pool: Option<&Pool>,
4499) {
4500    let gpr = cols / GROUP_SIZE;
4501    let p1 = SendMut(o1.as_mut_ptr());
4502    let p2 = SendMut(o2.as_mut_ptr());
4503    if a8w8_enabled() {
4504        let a1 = split_act(x1);
4505        let a2 = split_act(x2);
4506        let run = move |start: usize, end: usize| {
4507            for r in start..end {
4508                let mut v1 = dot_q4t_row_i8(bytes, r, gpr, &a1.xq) * a1.sx;
4509                let mut v2 = dot_q4t_row_i8(bytes, r, gpr, &a2.xq) * a2.sx;
4510                for &(j, xv) in &a1.outliers {
4511                    let (w, s) = q4t_outlier(bytes, r, gpr, j);
4512                    v1 += w * s * xv;
4513                }
4514                for &(j, xv) in &a2.outliers {
4515                    let (w, s) = q4t_outlier(bytes, r, gpr, j);
4516                    v2 += w * s * xv;
4517                }
4518                // SAFETY: disjoint row ranges per worker.
4519                unsafe {
4520                    *p1.at(r) = v1;
4521                    *p2.at(r) = v2;
4522                }
4523            }
4524        };
4525        dispatch_rows(pool, rows, &run);
4526        return;
4527    }
4528    let run = move |start: usize, end: usize| {
4529        for r in start..end {
4530            // SAFETY: disjoint row ranges per worker.
4531            unsafe {
4532                *p1.at(r) = q4t_row_exact(bytes, r, gpr, x1);
4533                *p2.at(r) = q4t_row_exact(bytes, r, gpr, x2);
4534            }
4535        }
4536    };
4537    dispatch_rows(pool, rows, &run);
4538}
4539
4540/// Batched q4_tiled matmat: each row's tiles stream once per microbatch.
4541#[allow(clippy::too_many_arguments)]
4542/// Prefill GEMM through Accelerate for group-quantized codecs: a
4543/// caller-supplied row dequantizer fills f32 tiles (pool-parallel) and
4544/// each tile rides the AMX with one sgemm — the generic sibling of
4545/// `qmatmat_accel` (q8). Numerics are f32-GEMM (tolerance class);
4546/// decode (b=1) never takes this path.
4547#[cfg(target_os = "macos")]
4548fn dequant_matmat_accel(
4549    dequant_row: &(dyn Fn(usize, &mut [f32]) + Sync),
4550    xs_all: &[f32],
4551    b: usize,
4552    rows: usize,
4553    cols: usize,
4554    out: &mut [f32],
4555    pool: Option<&Pool>,
4556) {
4557    const TR: usize = 2048;
4558    thread_local! {
4559        static WTILE: std::cell::RefCell<Vec<f32>> = const { std::cell::RefCell::new(Vec::new()) };
4560    }
4561    WTILE.with(|wt| {
4562        let mut wtile = wt.borrow_mut();
4563        wtile.resize(TR * cols, 0.0);
4564        let mut r0 = 0usize;
4565        while r0 < rows {
4566            let tr = TR.min(rows - r0);
4567            let wt_addr = SendMut(wtile.as_mut_ptr());
4568            let run = |start: usize, end: usize| {
4569                for r in start..end {
4570                    // SAFETY: workers cover disjoint r ranges.
4571                    let dst = unsafe { std::slice::from_raw_parts_mut(wt_addr.at(r * cols), cols) };
4572                    dequant_row(r0 + r, dst);
4573                }
4574            };
4575            dispatch_rows(pool, tr, &run);
4576            unsafe {
4577                accel_blas::cblas_sgemm(
4578                    101, // RowMajor
4579                    111, // NoTrans A
4580                    112, // Trans B
4581                    b as i32,
4582                    tr as i32,
4583                    cols as i32,
4584                    1.0,
4585                    xs_all.as_ptr(),
4586                    cols as i32,
4587                    wtile.as_ptr(),
4588                    cols as i32,
4589                    0.0,
4590                    out.as_mut_ptr().add(r0),
4591                    rows as i32,
4592                );
4593            }
4594            r0 += tr;
4595        }
4596    });
4597}
4598
4599fn q4t_matmat(
4600    bytes: &[u8],
4601    xs_all: &[f32],
4602    b: usize,
4603    rows: usize,
4604    cols: usize,
4605    out: &mut [f32],
4606    pool: Option<&Pool>,
4607) {
4608    debug_assert_eq!(out.len(), b * rows);
4609    let gpr = cols / GROUP_SIZE;
4610    // Wide batches ride the AMX like q8's qmatmat: on Apple silicon
4611    // the dequant-tile sgemm is an order above the SDOT row loop for
4612    // prefill shapes (imagegen DiT forwards are exactly this).
4613    #[cfg(target_os = "macos")]
4614    if b >= 8 && rows * cols >= 500_000 && accel_gemm_enabled() {
4615        dequant_matmat_accel(
4616            &|r, dst| {
4617                for gi in 0..gpr {
4618                    let tile = &bytes[(r * gpr + gi) * Q4_TILE..(r * gpr + gi + 1) * Q4_TILE];
4619                    let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
4620                    for (k, &bb) in tile[2..].iter().enumerate() {
4621                        dst[gi * GROUP_SIZE + k * 2] = ((bb & 0x0F) as f32 - 8.0) * s;
4622                        dst[gi * GROUP_SIZE + k * 2 + 1] = (((bb >> 4) & 0x0F) as f32 - 8.0) * s;
4623                    }
4624                }
4625            },
4626            xs_all,
4627            b,
4628            rows,
4629            cols,
4630            out,
4631            pool,
4632        );
4633        return;
4634    }
4635    let out_addr = SendMut(out.as_mut_ptr());
4636    if a8w8_enabled() {
4637        let acts: Vec<SplitAct> = (0..b)
4638            .map(|bi| split_act(&xs_all[bi * cols..(bi + 1) * cols]))
4639            .collect();
4640        let acts = &acts;
4641        #[cfg(target_arch = "x86_64")]
4642        let blocked_ok = avx2_enabled()
4643            && std::env::var("CMF_X86_BLOCKED")
4644                .map(|v| v != "0")
4645                .unwrap_or(true);
4646        #[cfg(target_arch = "aarch64")]
4647        let blocked_ok = sdot_enabled()
4648            && std::env::var("CMF_X86_BLOCKED")
4649                .map(|v| v != "0")
4650                .unwrap_or(true);
4651        #[cfg(not(any(target_arch = "x86_64", target_arch = "aarch64")))]
4652        let blocked_ok = false;
4653        let run = move |start: usize, end: usize| {
4654            for r in start..end {
4655                let mut bi = 0usize;
4656                #[cfg(target_arch = "aarch64")]
4657                if blocked_ok {
4658                    while bi + 4 <= acts.len() {
4659                        let xs = [
4660                            acts[bi].xq.as_slice(),
4661                            acts[bi + 1].xq.as_slice(),
4662                            acts[bi + 2].xq.as_slice(),
4663                            acts[bi + 3].xq.as_slice(),
4664                        ];
4665                        let d = unsafe { dot_q4t_row_1x4_sdot(bytes, r, gpr, xs) };
4666                        for k in 0..4 {
4667                            let act = &acts[bi + k];
4668                            let mut acc = d[k] * act.sx;
4669                            for &(j, xv) in &act.outliers {
4670                                let (w, sc) = q4t_outlier(bytes, r, gpr, j);
4671                                acc += w * sc * xv;
4672                            }
4673                            // SAFETY: disjoint (bi, r) cells per worker.
4674                            unsafe { *out_addr.at((bi + k) * rows + r) = acc };
4675                        }
4676                        bi += 4;
4677                    }
4678                }
4679                #[cfg(target_arch = "x86_64")]
4680                if blocked_ok {
4681                    while bi + 4 <= acts.len() {
4682                        let xs = [
4683                            acts[bi].xq.as_slice(),
4684                            acts[bi + 1].xq.as_slice(),
4685                            acts[bi + 2].xq.as_slice(),
4686                            acts[bi + 3].xq.as_slice(),
4687                        ];
4688                        let d = unsafe {
4689                            if vnni_tiles_enabled() {
4690                                dot_q4t_row_1x4_vnni(bytes, r, gpr, xs)
4691                            } else {
4692                                dot_q4t_row_1x4_avx2(bytes, r, gpr, xs)
4693                            }
4694                        };
4695                        for k in 0..4 {
4696                            let act = &acts[bi + k];
4697                            let mut acc = d[k] * act.sx;
4698                            for &(j, xv) in &act.outliers {
4699                                let (w, sc) = q4t_outlier(bytes, r, gpr, j);
4700                                acc += w * sc * xv;
4701                            }
4702                            // SAFETY: disjoint (bi, r) cells per worker.
4703                            unsafe { *out_addr.at((bi + k) * rows + r) = acc };
4704                        }
4705                        bi += 4;
4706                    }
4707                }
4708                let _ = blocked_ok;
4709                while bi < acts.len() {
4710                    let act = &acts[bi];
4711                    let mut acc = dot_q4t_row_i8(bytes, r, gpr, &act.xq) * act.sx;
4712                    for &(j, xv) in &act.outliers {
4713                        let (w, s) = q4t_outlier(bytes, r, gpr, j);
4714                        acc += w * s * xv;
4715                    }
4716                    // SAFETY: disjoint (bi, r) cells per worker range.
4717                    unsafe { *out_addr.at(bi * rows + r) = acc };
4718                    bi += 1;
4719                }
4720            }
4721        };
4722        dispatch_rows(pool, rows, &run);
4723        return;
4724    }
4725    let run = move |start: usize, end: usize| {
4726        for r in start..end {
4727            for bi in 0..b {
4728                let x = &xs_all[bi * cols..(bi + 1) * cols];
4729                // SAFETY: disjoint (bi, r) cells per worker range.
4730                unsafe { *out_addr.at(bi * rows + r) = q4t_row_exact(bytes, r, gpr, x) };
4731            }
4732        }
4733    };
4734    dispatch_rows(pool, rows, &run);
4735}
4736
4737// ── q1 (dtype 12): binary weights, [f16 scale][4B sign bits] per
4738// 32-group tile. The kernel family mirrors q4_tiled: one sequential
4739// stream of 6-byte tiles, per-tile integer dot × scale, exact outlier
4740// correction (A8W8 contract), exact scalar path under CMF_SDOT=0. ──
4741
4742/// Per-32-group sums of the quantized activation — the ±1 identity's
4743/// shared half: `dot = −2·sdot(mask, x) − gsum[g]`, computed ONCE per
4744/// matvec and reused by every row.
4745fn q1_group_sums(xq: &[i8], gpr: usize) -> Vec<i32> {
4746    (0..gpr)
4747        .map(|gi| {
4748            xq[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE]
4749                .iter()
4750                .map(|&v| v as i32)
4751                .sum()
4752        })
4753        .collect()
4754}
4755
4756/// One q1 row via the A8W8 int8 path — mask-SDOT on ARM (no ±1
4757/// expansion at all), scalar bit loop elsewhere (AVX2 queued with the
4758/// x86 pass).
4759#[inline]
4760#[allow(unreachable_code)]
4761/// AVX2 q1 row via the same ±1 identity as the ARM sdot kernel: the
4762/// sign bits expand to a {0, −1} byte mask through shuffle+cmpeq, the
4763/// masked activation sums through maddubs(1, x&mask), and
4764/// `dot = −(2·masked_sum + Σx_group)` — bit-identical integer math.
4765#[cfg(target_arch = "x86_64")]
4766#[target_feature(enable = "avx2")]
4767unsafe fn dot_q1_row_avx2(bytes: &[u8], r: usize, gpr: usize, xq: &[i8], gsum: &[i32]) -> f32 {
4768    // SAFETY: callers uphold the 6B-tile and xq/gsum length contracts.
4769    unsafe {
4770        use core::arch::x86_64::*;
4771        // Byte j of the mask must replicate bits-byte j/8.
4772        let expand = _mm256_setr_epi8(
4773            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,
4774            3, 3, 3,
4775        );
4776        let bitsel = _mm256_setr_epi8(
4777            1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64,
4778            -128, 1, 2, 4, 8, 16, 32, 64, -128,
4779        );
4780        let ones8 = _mm256_set1_epi8(1);
4781        let ones16 = _mm256_set1_epi16(1);
4782        let mut acc = 0f32;
4783        for gi in 0..gpr {
4784            let t = bytes.as_ptr().add((r * gpr + gi) * Q1_TILE);
4785            let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
4786            let bits = u32::from_le_bytes([*t.add(2), *t.add(3), *t.add(4), *t.add(5)]);
4787            let bc = _mm256_shuffle_epi8(_mm256_set1_epi32(bits as i32), expand);
4788            let mask = _mm256_cmpeq_epi8(_mm256_and_si256(bc, bitsel), bitsel);
4789            let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
4790            let sel = _mm256_and_si256(x, mask);
4791            // Σ of selected i8 lanes: maddubs(1u8, sel_i8) pairs → madd.
4792            let p16 = _mm256_maddubs_epi16(ones8, sel);
4793            let d32 = _mm256_madd_epi16(p16, ones16);
4794            let hi128 = _mm256_extracti128_si256::<1>(d32);
4795            let s128 = _mm_add_epi32(_mm256_castsi256_si128(d32), hi128);
4796            let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
4797            let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
4798            let msum = _mm_cvtsi128_si32(s32);
4799            // The and-select keeps x UN-negated (unlike ARM's −1-mask
4800            // sdot): d = Σ_set − Σ_unset = 2·Σ_set − Σ_all.
4801            let d = 2 * msum - gsum[gi];
4802            acc += d as f32 * s;
4803        }
4804        acc
4805    }
4806}
4807
4808/// VNNI twin of `dot_q1_row_avx2`: the masked-select sum goes through
4809/// one `vpdpbusd(1u8, sel)` (see `dpbusd_hsum` — bit-identical).
4810#[cfg(target_arch = "x86_64")]
4811#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
4812unsafe fn dot_q1_row_vnni(bytes: &[u8], r: usize, gpr: usize, xq: &[i8], gsum: &[i32]) -> f32 {
4813    // SAFETY: callers uphold the 6B-tile and xq/gsum length contracts.
4814    unsafe {
4815        use core::arch::x86_64::*;
4816        let expand = _mm256_setr_epi8(
4817            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,
4818            3, 3, 3,
4819        );
4820        let bitsel = _mm256_setr_epi8(
4821            1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64,
4822            -128, 1, 2, 4, 8, 16, 32, 64, -128,
4823        );
4824        let ones8 = _mm256_set1_epi8(1);
4825        let mut acc = 0f32;
4826        for gi in 0..gpr {
4827            let t = bytes.as_ptr().add((r * gpr + gi) * Q1_TILE);
4828            let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
4829            let bits = u32::from_le_bytes([*t.add(2), *t.add(3), *t.add(4), *t.add(5)]);
4830            let bc = _mm256_shuffle_epi8(_mm256_set1_epi32(bits as i32), expand);
4831            let mask = _mm256_cmpeq_epi8(_mm256_and_si256(bc, bitsel), bitsel);
4832            let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
4833            let msum = dpbusd_hsum(ones8, _mm256_and_si256(x, mask));
4834            let d = 2 * msum - gsum[gi];
4835            acc += d as f32 * s;
4836        }
4837        acc
4838    }
4839}
4840
4841/// VNNI twin of `dot_q1_row_1x4_avx2` (see `dpbusd_hsum`).
4842#[cfg(target_arch = "x86_64")]
4843#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
4844unsafe fn dot_q1_row_1x4_vnni(
4845    bytes: &[u8],
4846    r: usize,
4847    gpr: usize,
4848    xs: [&[i8]; 4],
4849    gsums: [&[i32]; 4],
4850) -> [f32; 4] {
4851    // SAFETY: callers uphold the 6B-tile and xq/gsum length contracts.
4852    unsafe {
4853        use core::arch::x86_64::*;
4854        let expand = _mm256_setr_epi8(
4855            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,
4856            3, 3, 3,
4857        );
4858        let bitsel = _mm256_setr_epi8(
4859            1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64,
4860            -128, 1, 2, 4, 8, 16, 32, 64, -128,
4861        );
4862        let ones8 = _mm256_set1_epi8(1);
4863        let mut acc = [0f32; 4];
4864        for gi in 0..gpr {
4865            let t = bytes.as_ptr().add((r * gpr + gi) * Q1_TILE);
4866            let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
4867            let bits = u32::from_le_bytes([*t.add(2), *t.add(3), *t.add(4), *t.add(5)]);
4868            let bc = _mm256_shuffle_epi8(_mm256_set1_epi32(bits as i32), expand);
4869            let mask = _mm256_cmpeq_epi8(_mm256_and_si256(bc, bitsel), bitsel);
4870            for (k, xq) in xs.iter().enumerate() {
4871                let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
4872                let msum = dpbusd_hsum(ones8, _mm256_and_si256(x, mask));
4873                let d = 2 * msum - gsums[k][gi];
4874                acc[k] += d as f32 * s;
4875            }
4876        }
4877        acc
4878    }
4879}
4880
4881/// The blocked 1×4 flavor: the expanded bit mask serves four activation
4882/// streams per group (mask build once, four select+reduce chains).
4883#[cfg(target_arch = "x86_64")]
4884#[target_feature(enable = "avx2")]
4885unsafe fn dot_q1_row_1x4_avx2(
4886    bytes: &[u8],
4887    r: usize,
4888    gpr: usize,
4889    xs: [&[i8]; 4],
4890    gsums: [&[i32]; 4],
4891) -> [f32; 4] {
4892    // SAFETY: callers uphold the 6B-tile and xq/gsum length contracts.
4893    unsafe {
4894        use core::arch::x86_64::*;
4895        let expand = _mm256_setr_epi8(
4896            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,
4897            3, 3, 3,
4898        );
4899        let bitsel = _mm256_setr_epi8(
4900            1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64,
4901            -128, 1, 2, 4, 8, 16, 32, 64, -128,
4902        );
4903        let ones8 = _mm256_set1_epi8(1);
4904        let ones16 = _mm256_set1_epi16(1);
4905        let mut acc = [0f32; 4];
4906        for gi in 0..gpr {
4907            let t = bytes.as_ptr().add((r * gpr + gi) * Q1_TILE);
4908            let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
4909            let bits = u32::from_le_bytes([*t.add(2), *t.add(3), *t.add(4), *t.add(5)]);
4910            let bc = _mm256_shuffle_epi8(_mm256_set1_epi32(bits as i32), expand);
4911            let mask = _mm256_cmpeq_epi8(_mm256_and_si256(bc, bitsel), bitsel);
4912            for (k, xq) in xs.iter().enumerate() {
4913                let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
4914                let sel = _mm256_and_si256(x, mask);
4915                let p16 = _mm256_maddubs_epi16(ones8, sel);
4916                let d32 = _mm256_madd_epi16(p16, ones16);
4917                let hi128 = _mm256_extracti128_si256::<1>(d32);
4918                let s128 = _mm_add_epi32(_mm256_castsi256_si128(d32), hi128);
4919                let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
4920                let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
4921                let msum = _mm_cvtsi128_si32(s32);
4922                let d = 2 * msum - gsums[k][gi];
4923                acc[k] += d as f32 * s;
4924            }
4925        }
4926        acc
4927    }
4928}
4929
4930#[allow(unreachable_code)]
4931fn dot_q1_row_i8(bytes: &[u8], r: usize, gpr: usize, xq: &[i8], gsum: &[i32]) -> f32 {
4932    #[cfg(target_arch = "aarch64")]
4933    unsafe {
4934        return dot_q1_row_sdot(bytes, r, gpr, xq, gsum);
4935    }
4936    #[cfg(target_arch = "x86_64")]
4937    if avx2_enabled() {
4938        unsafe {
4939            if vnni_tiles_enabled() {
4940                return dot_q1_row_vnni(bytes, r, gpr, xq, gsum);
4941            }
4942            return dot_q1_row_avx2(bytes, r, gpr, xq, gsum);
4943        }
4944    }
4945    let _ = gsum;
4946    let mut acc = 0f32;
4947    for gi in 0..gpr {
4948        let tile = &bytes[(r * gpr + gi) * Q1_TILE..(r * gpr + gi + 1) * Q1_TILE];
4949        let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
4950        let mut d = 0i32;
4951        for (j, &b) in tile[2..].iter().enumerate() {
4952            for k in 0..8 {
4953                let w = ((b >> k) & 1) as i32 * 2 - 1;
4954                d += w * xq[gi * GROUP_SIZE + j * 8 + k] as i32;
4955            }
4956        }
4957        acc += d as f32 * s;
4958    }
4959    acc
4960}
4961
4962/// SDOT q1 row via the ±1 identity: the vtst mask (0xFF where the bit
4963/// is set, i.e. −1 as i8) feeds `sdot` DIRECTLY — no expansion to ±1
4964/// lanes at all — and `dot = −(2·sdot(mask, x) + Σx_group)`, with the
4965/// per-group activation sums shared across every row of the matvec.
4966/// Four tiles (128 weights) per iteration: integer dots reduce through
4967/// a vpaddq tree into ONE i32x4 that meets its four scales in a single
4968/// fused f32 multiply-add. Integer math throughout — bit-identical to
4969/// the scalar ±1 reference.
4970#[cfg(target_arch = "aarch64")]
4971#[target_feature(enable = "neon,dotprod")]
4972unsafe fn dot_q1_row_sdot(bytes: &[u8], r: usize, gpr: usize, xq: &[i8], gsum: &[i32]) -> f32 {
4973    // SAFETY: callers uphold slice-length contracts (6B tile per group,
4974    // xq.len() == gpr·GROUP_SIZE, gsum.len() == gpr).
4975    unsafe {
4976        use core::arch::aarch64::*;
4977        use core::arch::asm;
4978        const MASKS: [u8; 16] = [1, 2, 4, 8, 16, 32, 64, 128, 1, 2, 4, 8, 16, 32, 64, 128];
4979        let m = vld1q_u8(MASKS.as_ptr());
4980        // One tile's −Σ_set(x) as an UNREDUCED i32x4 (two mask-sdots).
4981        macro_rules! tile_dot {
4982            ($t:expr, $x:expr) => {{
4983                let v0 = vcombine_u8(vdup_n_u8(*$t.add(2)), vdup_n_u8(*$t.add(3)));
4984                let v1 = vcombine_u8(vdup_n_u8(*$t.add(4)), vdup_n_u8(*$t.add(5)));
4985                let w0 = vreinterpretq_s8_u8(vtstq_u8(v0, m));
4986                let w1 = vreinterpretq_s8_u8(vtstq_u8(v1, m));
4987                let x0 = vld1q_s8($x);
4988                let x1 = vld1q_s8($x.add(16));
4989                let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
4990                asm!(
4991                    "sdot {a0:v}.4s, {w0:v}.16b, {x0:v}.16b",
4992                    "sdot {a1:v}.4s, {w1:v}.16b, {x1:v}.16b",
4993                    a0 = inout(vreg) a0, a1 = inout(vreg) a1,
4994                    w0 = in(vreg) w0, x0 = in(vreg) x0, w1 = in(vreg) w1, x1 = in(vreg) x1,
4995                    options(pure, nomem, nostack),
4996                );
4997                vaddq_s32(a0, a1)
4998            }};
4999        }
5000        // TBL unpack over PAIR loads: one vld1q covers two 6B tiles
5001        // ([s s b b b b][s s b b b b] + 4B slack), TBL replicates each
5002        // bit-byte across 8 lanes for vtst, and the four scales gather
5003        // through tbl2 into one fcvtl — the 16 ld1r broadcast loads and
5004        // 4 branchy software f16 conversions per 128 weights (the
5005        // measured load-port wall of this kernel) become 2 vector
5006        // loads + 9 table lookups. Integer math order is unchanged —
5007        // bit-identical results (FCVTL is exact on every f16).
5008        const IW00: [u8; 16] = [2, 2, 2, 2, 2, 2, 2, 2, 3, 3, 3, 3, 3, 3, 3, 3];
5009        const IW01: [u8; 16] = [4, 4, 4, 4, 4, 4, 4, 4, 5, 5, 5, 5, 5, 5, 5, 5];
5010        const IW10: [u8; 16] = [8, 8, 8, 8, 8, 8, 8, 8, 9, 9, 9, 9, 9, 9, 9, 9];
5011        const IW11: [u8; 16] = [
5012            10, 10, 10, 10, 10, 10, 10, 10, 11, 11, 11, 11, 11, 11, 11, 11,
5013        ];
5014        const ISC: [u8; 8] = [0, 1, 6, 7, 16, 17, 22, 23];
5015        let (iw00, iw01) = (vld1q_u8(IW00.as_ptr()), vld1q_u8(IW01.as_ptr()));
5016        let (iw10, iw11) = (vld1q_u8(IW10.as_ptr()), vld1q_u8(IW11.as_ptr()));
5017        let isc = vld1_u8(ISC.as_ptr());
5018        // One tile's −Σ_set(x) from a TBL-unpacked pair load.
5019        macro_rules! tile_dot_tbl {
5020            ($ld:expr, $i0:expr, $i1:expr, $x:expr) => {{
5021                let w0 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8($ld, $i0), m));
5022                let w1 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8($ld, $i1), m));
5023                let x0 = vld1q_s8($x);
5024                let x1 = vld1q_s8($x.add(16));
5025                let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
5026                asm!(
5027                    "sdot {a0:v}.4s, {w0:v}.16b, {x0:v}.16b",
5028                    "sdot {a1:v}.4s, {w1:v}.16b, {x1:v}.16b",
5029                    a0 = inout(vreg) a0, a1 = inout(vreg) a1,
5030                    w0 = in(vreg) w0, x0 = in(vreg) x0, w1 = in(vreg) w1, x1 = in(vreg) x1,
5031                    options(pure, nomem, nostack),
5032                );
5033                vaddq_s32(a0, a1)
5034            }};
5035        }
5036        let base = bytes.as_ptr().add(r * gpr * Q1_TILE);
5037        let row_base = r * gpr * Q1_TILE;
5038        let abs_end = bytes.len();
5039        let xp = xq.as_ptr();
5040        let gp = gsum.as_ptr();
5041        let mut accv = vdupq_n_f32(0.0);
5042        let mut gi = 0;
5043        // The second pair load reads 4B past tile gi+3 — stay inside
5044        // the payload slice (only the file's final tiles fall back).
5045        while gi + 4 <= gpr && row_base + (gi + 4) * Q1_TILE + 4 <= abs_end {
5046            let t0 = base.add(gi * Q1_TILE);
5047            let ld_a = vld1q_u8(t0);
5048            let ld_b = vld1q_u8(t0.add(2 * Q1_TILE));
5049            let d0 = tile_dot_tbl!(ld_a, iw00, iw01, xp.add(gi * GROUP_SIZE));
5050            let d1 = tile_dot_tbl!(ld_a, iw10, iw11, xp.add((gi + 1) * GROUP_SIZE));
5051            let d2 = tile_dot_tbl!(ld_b, iw00, iw01, xp.add((gi + 2) * GROUP_SIZE));
5052            let d3 = tile_dot_tbl!(ld_b, iw10, iw11, xp.add((gi + 3) * GROUP_SIZE));
5053            // [−Σ0, −Σ1, −Σ2, −Σ3] → dots = −(2·Σset_neg + gsum)
5054            let neg = vpaddq_s32(vpaddq_s32(d0, d1), vpaddq_s32(d2, d3));
5055            let g = vld1q_s32(gp.add(gi));
5056            let dots = vnegq_s32(vaddq_s32(vshlq_n_s32::<1>(neg), g));
5057            let sc16 = vqtbl2_u8(uint8x16x2_t(ld_a, ld_b), isc);
5058            let scf: float32x4_t;
5059            asm!(
5060                "fcvtl {o:v}.4s, {i:v}.4h",
5061                o = out(vreg) scf, i = in(vreg) sc16,
5062                options(pure, nomem, nostack),
5063            );
5064            accv = vfmaq_f32(accv, vcvtq_f32_s32(dots), scf);
5065            gi += 4;
5066        }
5067        let mut acc = vaddvq_f32(accv);
5068        while gi < gpr {
5069            let t = base.add(gi * Q1_TILE);
5070            let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
5071            let d = vaddvq_s32(tile_dot!(t, xp.add(gi * GROUP_SIZE)));
5072            acc += (-(2 * d + *gp.add(gi))) as f32 * s;
5073            gi += 1;
5074        }
5075        acc
5076    }
5077}
5078
5079/// Blocked q1 1×4: one TBL unpack of the tile pair serves FOUR
5080/// activation streams (prefill amortization — the same idea as the
5081/// AVX2 twin; per stream the group order, fma order and tail match the
5082/// single-row kernel exactly, so batch == matvec bit-for-bit).
5083#[cfg(target_arch = "aarch64")]
5084#[target_feature(enable = "neon,dotprod")]
5085unsafe fn dot_q1_row_1x4_sdot(
5086    bytes: &[u8],
5087    r: usize,
5088    gpr: usize,
5089    xs: [&[i8]; 4],
5090    gs: [&[i32]; 4],
5091) -> [f32; 4] {
5092    // SAFETY: same slice-length contracts as `dot_q1_row_sdot`, ×4.
5093    unsafe {
5094        use core::arch::aarch64::*;
5095        use core::arch::asm;
5096        const MASKS: [u8; 16] = [1, 2, 4, 8, 16, 32, 64, 128, 1, 2, 4, 8, 16, 32, 64, 128];
5097        const IW00: [u8; 16] = [2, 2, 2, 2, 2, 2, 2, 2, 3, 3, 3, 3, 3, 3, 3, 3];
5098        const IW01: [u8; 16] = [4, 4, 4, 4, 4, 4, 4, 4, 5, 5, 5, 5, 5, 5, 5, 5];
5099        const IW10: [u8; 16] = [8, 8, 8, 8, 8, 8, 8, 8, 9, 9, 9, 9, 9, 9, 9, 9];
5100        const IW11: [u8; 16] = [
5101            10, 10, 10, 10, 10, 10, 10, 10, 11, 11, 11, 11, 11, 11, 11, 11,
5102        ];
5103        const ISC: [u8; 8] = [0, 1, 6, 7, 16, 17, 22, 23];
5104        let m = vld1q_u8(MASKS.as_ptr());
5105        let (iw00, iw01) = (vld1q_u8(IW00.as_ptr()), vld1q_u8(IW01.as_ptr()));
5106        let (iw10, iw11) = (vld1q_u8(IW10.as_ptr()), vld1q_u8(IW11.as_ptr()));
5107        let isc = vld1_u8(ISC.as_ptr());
5108        macro_rules! sdot2 {
5109            ($w0:expr, $w1:expr, $x:expr) => {{
5110                let x0 = vld1q_s8($x);
5111                let x1 = vld1q_s8($x.add(16));
5112                let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
5113                asm!(
5114                    "sdot {a0:v}.4s, {w0:v}.16b, {x0:v}.16b",
5115                    "sdot {a1:v}.4s, {w1:v}.16b, {x1:v}.16b",
5116                    a0 = inout(vreg) a0, a1 = inout(vreg) a1,
5117                    w0 = in(vreg) $w0, x0 = in(vreg) x0, w1 = in(vreg) $w1, x1 = in(vreg) x1,
5118                    options(pure, nomem, nostack),
5119                );
5120                vaddq_s32(a0, a1)
5121            }};
5122        }
5123        let base = bytes.as_ptr().add(r * gpr * Q1_TILE);
5124        let row_base = r * gpr * Q1_TILE;
5125        let abs_end = bytes.len();
5126        let mut accv = [vdupq_n_f32(0.0); 4];
5127        let mut gi = 0;
5128        while gi + 4 <= gpr && row_base + (gi + 4) * Q1_TILE + 4 <= abs_end {
5129            let t0 = base.add(gi * Q1_TILE);
5130            let ld_a = vld1q_u8(t0);
5131            let ld_b = vld1q_u8(t0.add(2 * Q1_TILE));
5132            // Unpack ONCE — eight ±mask vectors serve all four streams.
5133            let w00 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8(ld_a, iw00), m));
5134            let w01 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8(ld_a, iw01), m));
5135            let w10 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8(ld_a, iw10), m));
5136            let w11 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8(ld_a, iw11), m));
5137            let w20 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8(ld_b, iw00), m));
5138            let w21 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8(ld_b, iw01), m));
5139            let w30 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8(ld_b, iw10), m));
5140            let w31 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8(ld_b, iw11), m));
5141            let sc16 = vqtbl2_u8(uint8x16x2_t(ld_a, ld_b), isc);
5142            let scf: float32x4_t;
5143            asm!(
5144                "fcvtl {o:v}.4s, {i:v}.4h",
5145                o = out(vreg) scf, i = in(vreg) sc16,
5146                options(pure, nomem, nostack),
5147            );
5148            for k in 0..4 {
5149                let xp = xs[k].as_ptr();
5150                let d0 = sdot2!(w00, w01, xp.add(gi * GROUP_SIZE));
5151                let d1 = sdot2!(w10, w11, xp.add((gi + 1) * GROUP_SIZE));
5152                let d2 = sdot2!(w20, w21, xp.add((gi + 2) * GROUP_SIZE));
5153                let d3 = sdot2!(w30, w31, xp.add((gi + 3) * GROUP_SIZE));
5154                let neg = vpaddq_s32(vpaddq_s32(d0, d1), vpaddq_s32(d2, d3));
5155                let g = vld1q_s32(gs[k].as_ptr().add(gi));
5156                let dots = vnegq_s32(vaddq_s32(vshlq_n_s32::<1>(neg), g));
5157                accv[k] = vfmaq_f32(accv[k], vcvtq_f32_s32(dots), scf);
5158            }
5159            gi += 4;
5160        }
5161        let mut acc = [
5162            vaddvq_f32(accv[0]),
5163            vaddvq_f32(accv[1]),
5164            vaddvq_f32(accv[2]),
5165            vaddvq_f32(accv[3]),
5166        ];
5167        while gi < gpr {
5168            let t = base.add(gi * Q1_TILE);
5169            let sc = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
5170            let v0 = vcombine_u8(vdup_n_u8(*t.add(2)), vdup_n_u8(*t.add(3)));
5171            let v1 = vcombine_u8(vdup_n_u8(*t.add(4)), vdup_n_u8(*t.add(5)));
5172            let w0 = vreinterpretq_s8_u8(vtstq_u8(v0, m));
5173            let w1 = vreinterpretq_s8_u8(vtstq_u8(v1, m));
5174            for k in 0..4 {
5175                let d = vaddvq_s32(sdot2!(w0, w1, xs[k].as_ptr().add(gi * GROUP_SIZE)));
5176                acc[k] += (-(2 * d + *gs[k].as_ptr().add(gi))) as f32 * sc;
5177            }
5178            gi += 1;
5179        }
5180        acc
5181    }
5182}
5183
5184/// (weight ±1, scale) of one q1 element — the exact outlier term.
5185#[inline]
5186fn q1_outlier(bytes: &[u8], r: usize, gpr: usize, j: usize) -> (f32, f32) {
5187    let gi = j / GROUP_SIZE;
5188    let k = j % GROUP_SIZE;
5189    let tile = &bytes[(r * gpr + gi) * Q1_TILE..(r * gpr + gi + 1) * Q1_TILE];
5190    let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
5191    let bit = (tile[2 + k / 8] >> (k % 8)) & 1;
5192    ((bit as i32 * 2 - 1) as f32, s)
5193}
5194
5195/// Exact scalar q1 row (CMF_SDOT=0 contract).
5196#[inline]
5197fn q1_row_exact(bytes: &[u8], r: usize, gpr: usize, x: &[f32]) -> f32 {
5198    let mut acc = 0f32;
5199    for gi in 0..gpr {
5200        let tile = &bytes[(r * gpr + gi) * Q1_TILE..(r * gpr + gi + 1) * Q1_TILE];
5201        let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
5202        let xg = &x[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
5203        let mut ga = 0f32;
5204        for (j, &b) in tile[2..].iter().enumerate() {
5205            for k in 0..8 {
5206                ga += (((b >> k) & 1) as f32 * 2.0 - 1.0) * xg[j * 8 + k];
5207            }
5208        }
5209        acc += ga * s;
5210    }
5211    acc
5212}
5213
5214/// One q1 row range via A8W8 (the body of `q1_matvec`'s hot loop,
5215/// extracted so multi-matrix jobs drive the same kernel).
5216#[allow(clippy::too_many_arguments)]
5217fn q1_range_a8w8(
5218    bytes: &[u8],
5219    gpr: usize,
5220    act: &SplitAct,
5221    gsum: &[i32],
5222    out: SendMut,
5223    start: usize,
5224    end: usize,
5225) {
5226    for r in start..end {
5227        let mut acc = dot_q1_row_i8(bytes, r, gpr, &act.xq, gsum) * act.sx;
5228        for &(j, xv) in &act.outliers {
5229            let (w, s) = q1_outlier(bytes, r, gpr, j);
5230            acc += w * s * xv;
5231        }
5232        // SAFETY: disjoint row ranges per worker.
5233        unsafe { *out.at(r) = acc };
5234    }
5235}
5236
5237/// Exact-scalar q1 row range (CMF_SDOT=0 contract).
5238fn q1_range_f32(bytes: &[u8], gpr: usize, x: &[f32], out: SendMut, start: usize, end: usize) {
5239    for r in start..end {
5240        // SAFETY: disjoint row ranges per worker.
5241        unsafe { *out.at(r) = q1_row_exact(bytes, r, gpr, x) };
5242    }
5243}
5244
5245/// q1t per-row overlay locator. After the base (`base_len`) come
5246/// `[u32 row_ptr[rows+1]]` then `[(u16 col, f16 val)]` grouped by row (row
5247/// `r`'s entries are `[row_ptr[r], row_ptr[r+1])`). Returns
5248/// `(row_ptr offset, entries offset, present)`.
5249fn q1t_overlay(bytes: &[u8], base_len: usize, rows: usize) -> (usize, usize, bool) {
5250    let entries = base_len + (rows + 1) * 4;
5251    (base_len, entries, entries <= bytes.len())
5252}
5253
5254/// Read `row_ptr[r]` from the overlay's prefix-sum table.
5255#[inline]
5256fn q1t_rowptr(bytes: &[u8], rp_off: usize, r: usize) -> usize {
5257    let o = rp_off + r * 4;
5258    u32::from_le_bytes([bytes[o], bytes[o + 1], bytes[o + 2], bytes[o + 3]]) as usize
5259}
5260
5261/// Byte → the 5 ternary signs it packs `{−1,0,+1}` as f32, precomputed so
5262/// decoding a q1t code is a table load, not the base-3 divide/modulo per
5263/// weight (division is ~20–40× the cost of a load). Built at compile time.
5264const SIGN5: [[f32; 5]; 256] = {
5265    let mut lut = [[0.0f32; 5]; 256];
5266    let pow3 = [1u16, 3, 9, 27, 81];
5267    let mut byte = 0usize;
5268    while byte < 256 {
5269        let mut i = 0usize;
5270        while i < 5 {
5271            let code = (byte as u16 / pow3[i]) % 3;
5272            lut[byte][i] = if code == 1 {
5273                1.0
5274            } else if code == 2 {
5275                -1.0
5276            } else {
5277                0.0
5278            };
5279            i += 1;
5280        }
5281        byte += 1;
5282    }
5283    lut
5284};
5285
5286/// Same table, as i8 signs — the operand for the int8 SDOT base kernel.
5287const SIGN5_I8: [[i8; 5]; 256] = {
5288    let mut lut = [[0i8; 5]; 256];
5289    let pow3 = [1u16, 3, 9, 27, 81];
5290    let mut byte = 0usize;
5291    while byte < 256 {
5292        let mut i = 0usize;
5293        while i < 5 {
5294            let code = (byte as u16 / pow3[i]) % 3;
5295            lut[byte][i] = if code == 1 {
5296                1
5297            } else if code == 2 {
5298                -1
5299            } else {
5300                0
5301            };
5302            i += 1;
5303        }
5304        byte += 1;
5305    }
5306    lut
5307};
5308
5309/// The same 5 i8 signs packed into a u64 (`[s0 s1 s2 s3 s4 0 0 0]`, LE) so the
5310/// group unpack is 7 unaligned u64 stores at offsets 0,5,10,…,30 instead of
5311/// six 5-byte copies + LUT indexing — each store's trailing zeros are fixed by
5312/// the next store, and the last one runs 6 B past the 32nd weight (the unpack
5313/// buffer is padded to 40). This is the decode/prefill hot inner op.
5314const SIGN5_U64: [u64; 256] = {
5315    let mut lut = [0u64; 256];
5316    let pow3 = [1u16, 3, 9, 27, 81];
5317    let mut byte = 0usize;
5318    while byte < 256 {
5319        let mut v = 0u64;
5320        let mut i = 0usize;
5321        while i < 5 {
5322            let code = (byte as u16 / pow3[i]) % 3;
5323            let s: u8 = if code == 1 {
5324                1
5325            } else if code == 2 {
5326                0xFF
5327            } else {
5328                0
5329            };
5330            v |= (s as u64) << (i * 8);
5331            i += 1;
5332        }
5333        lut[byte] = v;
5334        byte += 1;
5335    }
5336    lut
5337};
5338
5339/// Ternary base weight at `(row r, col j)` = `sign(code)·s_group`. Used to add
5340/// back activation-outlier columns, whose `x` was zeroed for the int8 bulk dot
5341/// (`split_act`). At a weight-outlier position the code is 0, so this is 0 and
5342/// the overlay correction owns that column — no double counting.
5343#[inline]
5344fn q1t_base_weight(bytes: &[u8], r: usize, gpr: usize, j: usize) -> f32 {
5345    const TILE: usize = cortiq_core::quant::Q1T_TILE;
5346    let off = (r * gpr + j / GROUP_SIZE) * TILE;
5347    let s = f16_to_f32(u16::from_le_bytes([bytes[off], bytes[off + 1]]));
5348    let within = j % GROUP_SIZE;
5349    SIGN5[bytes[off + 2 + within / 5] as usize][within % 5] * s
5350}
5351
5352/// One 32-group int8 dot via two SDOTs. Bit-exact vs the scalar i8 sum
5353/// (integer accumulation is order-independent).
5354#[cfg(target_arch = "aarch64")]
5355#[target_feature(enable = "neon,dotprod")]
5356#[inline]
5357unsafe fn sdot32_i8(w: *const i8, x: *const i8) -> i32 {
5358    // SAFETY: caller guarantees 32 readable i8 at each pointer.
5359    unsafe {
5360        use core::arch::aarch64::*;
5361        use core::arch::asm;
5362        let w0 = vld1q_s8(w);
5363        let w1 = vld1q_s8(w.add(16));
5364        let x0 = vld1q_s8(x);
5365        let x1 = vld1q_s8(x.add(16));
5366        let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
5367        asm!(
5368            "sdot {a0:v}.4s, {w0:v}.16b, {x0:v}.16b",
5369            "sdot {a1:v}.4s, {w1:v}.16b, {x1:v}.16b",
5370            a0 = inout(vreg) a0, a1 = inout(vreg) a1,
5371            w0 = in(vreg) w0, x0 = in(vreg) x0, w1 = in(vreg) w1, x1 = in(vreg) x1,
5372            options(pure, nomem, nostack),
5373        );
5374        vaddvq_s32(vaddq_s32(a0, a1))
5375    }
5376}
5377
5378/// One 32-group int8 dot via AVX2: signed·signed as `maddubs(|w|, sign(x,w))`
5379/// then `madd` and a horizontal reduce (the same idiom as `dot_q4t_row_avx2`).
5380#[cfg(target_arch = "x86_64")]
5381#[target_feature(enable = "avx2")]
5382#[inline]
5383unsafe fn i8dot32_avx2(w: *const i8, x: *const i8) -> i32 {
5384    // SAFETY: caller guarantees 32 readable i8 at each pointer.
5385    unsafe {
5386        use core::arch::x86_64::*;
5387        let wv = _mm256_loadu_si256(w as *const __m256i);
5388        let xv = _mm256_loadu_si256(x as *const __m256i);
5389        let p16 = _mm256_maddubs_epi16(_mm256_abs_epi8(wv), _mm256_sign_epi8(xv, wv));
5390        let d = _mm256_madd_epi16(p16, _mm256_set1_epi16(1));
5391        let hi128 = _mm256_extracti128_si256::<1>(d);
5392        let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
5393        let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
5394        let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
5395        _mm_cvtsi128_si32(s32)
5396    }
5397}
5398
5399/// Unpack one q1t group's base-3 codes into 32 i8 signs via 7 unaligned u64
5400/// stores (see `SIGN5_U64`). `dst` MUST have ≥ 40 bytes: the 7th store writes
5401/// `dst[30..38]`. Stores go in order so each one's trailing zeros are
5402/// overwritten by the next; the final 6 padding bytes are unused by the dot.
5403#[inline]
5404fn q1t_unpack_group_i8(codes: *const u8, dst: &mut [i8]) {
5405    debug_assert!(dst.len() >= 40);
5406    // SAFETY: codes points at 7 readable bytes; dst has ≥ 40 bytes so every
5407    // 8-byte store at offset bi*5 (bi ≤ 6 → ≤ 30) stays in bounds.
5408    unsafe {
5409        let p = dst.as_mut_ptr();
5410        for bi in 0..7 {
5411            core::ptr::write_unaligned(
5412                p.add(bi * 5) as *mut u64,
5413                SIGN5_U64[*codes.add(bi) as usize],
5414            );
5415        }
5416    }
5417}
5418
5419/// One 32-group int8 dot, arch-dispatched (the matmat inner loop, where the
5420/// row's signs are unpacked once and dotted against every batch input).
5421/// Callers are gated by `a8w8_enabled()`, so the target-feature arms are
5422/// reachable; the scalar arm is a non-SIMD-arch fallback.
5423#[inline]
5424fn q1t_i8dot32(w: *const i8, x: *const i8) -> i32 {
5425    #[cfg(target_arch = "aarch64")]
5426    unsafe {
5427        return sdot32_i8(w, x);
5428    }
5429    #[cfg(target_arch = "x86_64")]
5430    unsafe {
5431        return i8dot32_avx2(w, x);
5432    }
5433    #[allow(unreachable_code)]
5434    unsafe {
5435        let mut s = 0i32;
5436        for k in 0..GROUP_SIZE {
5437            s += *w.add(k) as i32 * *x.add(k) as i32;
5438        }
5439        s
5440    }
5441}
5442
5443#[inline]
5444unsafe fn q1t_unpack_reg_u64s(codes: *const u8) -> (u64, u64, u64, u64) {
5445    let (s0, s1, s2, s3, s4, s5, s6) = unsafe {
5446        (
5447            SIGN5_U64[*codes as usize],
5448            SIGN5_U64[*codes.add(1) as usize],
5449            SIGN5_U64[*codes.add(2) as usize],
5450            SIGN5_U64[*codes.add(3) as usize],
5451            SIGN5_U64[*codes.add(4) as usize],
5452            SIGN5_U64[*codes.add(5) as usize],
5453            SIGN5_U64[*codes.add(6) as usize],
5454        )
5455    };
5456
5457    let u0 = s0 | (s1 << 40);
5458    let u1 = (s1 >> 24) | (s2 << 16) | (s3 << 56);
5459    let u2 = (s3 >> 8) | (s4 << 32);
5460    let u3 = (s4 >> 32) | (s5 << 8) | (s6 << 48);
5461
5462    (u0, u1, u2, u3)
5463}
5464
5465/// One q1t row's int8 base dot: `Σ_group s·dot(signs, xq)` (before the shared
5466/// `sx`). Direct register unpacking (zero stack stores/loads, no STLF stalls).
5467/// ARM SDOT.
5468#[cfg(target_arch = "aarch64")]
5469#[target_feature(enable = "neon,dotprod")]
5470unsafe fn q1t_dot_row_sdot(bytes: &[u8], r: usize, gpr: usize, xq: &[i8]) -> f32 {
5471    use core::arch::aarch64::*;
5472    use core::arch::asm;
5473    unsafe {
5474        const TILE: usize = cortiq_core::quant::Q1T_TILE;
5475        let mut acc = 0f32;
5476        let bytes_ptr = bytes.as_ptr();
5477        let xq_ptr = xq.as_ptr();
5478        let row_off = r * gpr * TILE;
5479
5480        let gpr2 = gpr & !1;
5481        let mut gi = 0;
5482        while gi < gpr2 {
5483            let off0 = row_off + gi * TILE;
5484            let off1 = off0 + TILE;
5485            let s0 = f16_to_f32(u16::from_le_bytes([
5486                *bytes_ptr.add(off0),
5487                *bytes_ptr.add(off0 + 1),
5488            ]));
5489            let s1 = f16_to_f32(u16::from_le_bytes([
5490                *bytes_ptr.add(off1),
5491                *bytes_ptr.add(off1 + 1),
5492            ]));
5493
5494            let (u0_0, u1_0, u2_0, u3_0) = q1t_unpack_reg_u64s(bytes_ptr.add(off0 + 2));
5495            let (u0_1, u1_1, u2_1, u3_1) = q1t_unpack_reg_u64s(bytes_ptr.add(off1 + 2));
5496
5497            let w0_0 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u0_0), vcreate_u64(u1_0)));
5498            let w1_0 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u2_0), vcreate_u64(u3_0)));
5499            let w0_1 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u0_1), vcreate_u64(u1_1)));
5500            let w1_1 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u2_1), vcreate_u64(u3_1)));
5501
5502            let x0_0 = vld1q_s8(xq_ptr.add(gi * GROUP_SIZE));
5503            let x1_0 = vld1q_s8(xq_ptr.add(gi * GROUP_SIZE + 16));
5504            let x0_1 = vld1q_s8(xq_ptr.add((gi + 1) * GROUP_SIZE));
5505            let x1_1 = vld1q_s8(xq_ptr.add((gi + 1) * GROUP_SIZE + 16));
5506
5507            let (mut a0_0, mut a1_0) = (vdupq_n_s32(0), vdupq_n_s32(0));
5508            let (mut a0_1, mut a1_1) = (vdupq_n_s32(0), vdupq_n_s32(0));
5509            asm!(
5510                "sdot {a0_0:v}.4s, {w0_0:v}.16b, {x0_0:v}.16b",
5511                "sdot {a1_0:v}.4s, {w1_0:v}.16b, {x1_0:v}.16b",
5512                "sdot {a0_1:v}.4s, {w0_1:v}.16b, {x0_1:v}.16b",
5513                "sdot {a1_1:v}.4s, {w1_1:v}.16b, {x1_1:v}.16b",
5514                a0_0 = inout(vreg) a0_0, a1_0 = inout(vreg) a1_0,
5515                a0_1 = inout(vreg) a0_1, a1_1 = inout(vreg) a1_1,
5516                w0_0 = in(vreg) w0_0, x0_0 = in(vreg) x0_0, w1_0 = in(vreg) w1_0, x1_0 = in(vreg) x1_0,
5517                w0_1 = in(vreg) w0_1, x0_1 = in(vreg) x0_1, w1_1 = in(vreg) w1_1, x1_1 = in(vreg) x1_1,
5518                options(pure, nomem, nostack),
5519            );
5520            let d0 = vaddvq_s32(vaddq_s32(a0_0, a1_0));
5521            let d1 = vaddvq_s32(vaddq_s32(a0_1, a1_1));
5522            acc += d0 as f32 * s0 + d1 as f32 * s1;
5523            gi += 2;
5524        }
5525
5526        if gi < gpr {
5527            let off = row_off + gi * TILE;
5528            let s = f16_to_f32(u16::from_le_bytes([
5529                *bytes_ptr.add(off),
5530                *bytes_ptr.add(off + 1),
5531            ]));
5532            let (u0, u1, u2, u3) = q1t_unpack_reg_u64s(bytes_ptr.add(off + 2));
5533            let w0 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u0), vcreate_u64(u1)));
5534            let w1 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u2), vcreate_u64(u3)));
5535            let x0 = vld1q_s8(xq_ptr.add(gi * GROUP_SIZE));
5536            let x1 = vld1q_s8(xq_ptr.add(gi * GROUP_SIZE + 16));
5537            let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
5538            asm!(
5539                "sdot {a0:v}.4s, {w0:v}.16b, {x0:v}.16b",
5540                "sdot {a1:v}.4s, {w1:v}.16b, {x1:v}.16b",
5541                a0 = inout(vreg) a0, a1 = inout(vreg) a1,
5542                w0 = in(vreg) w0, x0 = in(vreg) x0, w1 = in(vreg) w1, x1 = in(vreg) x1,
5543                options(pure, nomem, nostack),
5544            );
5545            let d = vaddvq_s32(vaddq_s32(a0, a1));
5546            acc += d as f32 * s;
5547        }
5548        acc
5549    }
5550}
5551
5552/// x86 AVX2 mirror of `q1t_dot_row_sdot` (maddubs int8 dot per group).
5553#[cfg(target_arch = "x86_64")]
5554#[target_feature(enable = "avx2")]
5555unsafe fn q1t_dot_row_avx2(bytes: &[u8], r: usize, gpr: usize, xq: &[i8]) -> f32 {
5556    use core::arch::x86_64::*;
5557    unsafe {
5558        const TILE: usize = cortiq_core::quant::Q1T_TILE;
5559        let mut acc = 0f32;
5560        let bytes_ptr = bytes.as_ptr();
5561        let xq_ptr = xq.as_ptr();
5562        let row_off = r * gpr * TILE;
5563
5564        let ones = _mm256_set1_epi16(1);
5565        for gi in 0..gpr {
5566            let off = row_off + gi * TILE;
5567            let s = f16_to_f32(u16::from_le_bytes([
5568                *bytes_ptr.add(off),
5569                *bytes_ptr.add(off + 1),
5570            ]));
5571            let (u0, u1, u2, u3) = q1t_unpack_reg_u64s(bytes_ptr.add(off + 2));
5572            let wv = _mm256_set_epi64x(u3 as i64, u2 as i64, u1 as i64, u0 as i64);
5573            let xv = _mm256_loadu_si256(xq_ptr.add(gi * GROUP_SIZE) as *const __m256i);
5574            let p16 = _mm256_maddubs_epi16(_mm256_abs_epi8(wv), _mm256_sign_epi8(xv, wv));
5575            let d256 = _mm256_madd_epi16(p16, ones);
5576            let d128 = _mm_add_epi32(
5577                _mm256_castsi256_si128(d256),
5578                _mm256_extracti128_si256(d256, 1),
5579            );
5580            let d64 = _mm_add_epi32(d128, _mm_shuffle_epi32(d128, 0xee));
5581            let d32 = _mm_cvtsi128_si32(_mm_add_epi32(d64, _mm_shuffle_epi32(d64, 0x55)));
5582            acc += d32 as f32 * s;
5583        }
5584        acc
5585    }
5586}
5587
5588/// VNNI twin of `q1t_dot_row_avx2` (see `dpbusd_hsum`).
5589#[cfg(target_arch = "x86_64")]
5590#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
5591unsafe fn q1t_dot_row_vnni(bytes: &[u8], r: usize, gpr: usize, xq: &[i8]) -> f32 {
5592    use core::arch::x86_64::*;
5593    // SAFETY: same tile/xq contracts as `q1t_dot_row_avx2`.
5594    unsafe {
5595        const TILE: usize = cortiq_core::quant::Q1T_TILE;
5596        let mut acc = 0f32;
5597        let bytes_ptr = bytes.as_ptr();
5598        let xq_ptr = xq.as_ptr();
5599        let row_off = r * gpr * TILE;
5600        for gi in 0..gpr {
5601            let off = row_off + gi * TILE;
5602            let s = f16_to_f32(u16::from_le_bytes([
5603                *bytes_ptr.add(off),
5604                *bytes_ptr.add(off + 1),
5605            ]));
5606            let (u0, u1, u2, u3) = q1t_unpack_reg_u64s(bytes_ptr.add(off + 2));
5607            let wv = _mm256_set_epi64x(u3 as i64, u2 as i64, u1 as i64, u0 as i64);
5608            let xv = _mm256_loadu_si256(xq_ptr.add(gi * GROUP_SIZE) as *const __m256i);
5609            let d = dpbusd_hsum(_mm256_abs_epi8(wv), _mm256_sign_epi8(xv, wv));
5610            acc += d as f32 * s;
5611        }
5612        acc
5613    }
5614}
5615
5616/// Per-row int8 base dot, dispatched once per row (matvec decode hot path).
5617/// Callers are gated by `a8w8_enabled()`, so the target-feature kernels are
5618/// reachable.
5619#[inline]
5620fn q1t_dot_row_i8(bytes: &[u8], r: usize, gpr: usize, xq: &[i8]) -> f32 {
5621    #[cfg(target_arch = "aarch64")]
5622    unsafe {
5623        return q1t_dot_row_sdot(bytes, r, gpr, xq);
5624    }
5625    #[cfg(target_arch = "x86_64")]
5626    unsafe {
5627        if vnni_tiles_enabled() {
5628            return q1t_dot_row_vnni(bytes, r, gpr, xq);
5629        }
5630        return q1t_dot_row_avx2(bytes, r, gpr, xq);
5631    }
5632    #[allow(unreachable_code)]
5633    {
5634        const TILE: usize = cortiq_core::quant::Q1T_TILE;
5635        let mut acc = 0f32;
5636        let mut sg = [0i8; GROUP_SIZE + 8]; // +8 slack for the u64-store unpack
5637        for gi in 0..gpr {
5638            let off = (r * gpr + gi) * TILE;
5639            let s = f16_to_f32(u16::from_le_bytes([bytes[off], bytes[off + 1]]));
5640            q1t_unpack_group_i8(bytes.as_ptr().wrapping_add(off + 2), &mut sg);
5641            let mut d = 0i32;
5642            for k in 0..GROUP_SIZE {
5643                d += sg[k] as i32 * xq[gi * GROUP_SIZE + k] as i32;
5644            }
5645            acc += d as f32 * s;
5646        }
5647        acc
5648    }
5649}
5650
5651/// Σ over a row's outliers of `value·x[col]` — the correction that adds the
5652/// overlay's exact weights on top of the base dot. INVARIANT: the encoder
5653/// writes ternary code 0 at every outlier position (`quantize_q1t`), so the
5654/// base contributes nothing there and this is a plain `value·x`, not
5655/// `(value − base)·x` — no scattered per-outlier scale read. Row `r`'s entries
5656/// are the contiguous slice `[row_ptr[r], row_ptr[r+1])`, so no binary search.
5657fn q1t_row_outlier_correction(
5658    bytes: &[u8],
5659    r: usize,
5660    rp_off: usize,
5661    entries_off: usize,
5662    has_ov: bool,
5663    x: &[f32],
5664) -> f32 {
5665    if !has_ov {
5666        return 0.0;
5667    }
5668    let (c0, c1) = (
5669        q1t_rowptr(bytes, rp_off, r),
5670        q1t_rowptr(bytes, rp_off, r + 1),
5671    );
5672    let mut corr = 0f32;
5673    for p in c0..c1 {
5674        let e = entries_off + p * 4;
5675        let col = u16::from_le_bytes([bytes[e], bytes[e + 1]]) as usize;
5676        let val = f16_to_f32(u16::from_le_bytes([bytes[e + 2], bytes[e + 3]]));
5677        corr += val * x[col];
5678    }
5679    corr
5680}
5681
5682/// Dequantize one q1t row into `buf[..cols]` via the sign LUT (no division),
5683/// then apply the row's outliers (its `[row_ptr[r], row_ptr[r+1])` slice).
5684/// Used by the batched (prefill) path where the decode amortizes over the batch.
5685fn q1t_dequant_row(
5686    bytes: &[u8],
5687    r: usize,
5688    gpr: usize,
5689    rp_off: usize,
5690    entries_off: usize,
5691    has_ov: bool,
5692    buf: &mut [f32],
5693) {
5694    const TILE: usize = cortiq_core::quant::Q1T_TILE;
5695    for g in 0..gpr {
5696        let off = (r * gpr + g) * TILE;
5697        let s = f16_to_f32(u16::from_le_bytes([bytes[off], bytes[off + 1]]));
5698        let codes = &bytes[off + 2..off + TILE];
5699        let bc = g * GROUP_SIZE;
5700        // 6 full bytes (30 codes) + a 7th byte holding the last 2.
5701        for bi in 0..6 {
5702            let lut = &SIGN5[codes[bi] as usize];
5703            let d = &mut buf[bc + bi * 5..bc + bi * 5 + 5];
5704            for i in 0..5 {
5705                d[i] = lut[i] * s;
5706            }
5707        }
5708        let lut = &SIGN5[codes[6] as usize];
5709        buf[bc + 30] = lut[0] * s;
5710        buf[bc + 31] = lut[1] * s;
5711    }
5712    if !has_ov {
5713        return;
5714    }
5715    let (c0, c1) = (
5716        q1t_rowptr(bytes, rp_off, r),
5717        q1t_rowptr(bytes, rp_off, r + 1),
5718    );
5719    for p in c0..c1 {
5720        let e = entries_off + p * 4;
5721        let col = u16::from_le_bytes([bytes[e], bytes[e + 1]]) as usize;
5722        buf[col] = f16_to_f32(u16::from_le_bytes([bytes[e + 2], bytes[e + 3]]));
5723    }
5724}
5725
5726/// Add the sparse outlier overlay onto a base dot already in `out` (the GPU
5727/// computes the ternary base; the overlay stays on the CPU — its entries are
5728/// few and its per-row gather doesn't vectorize on the GPU). Row-parallel.
5729fn q1t_add_overlay(
5730    bytes: &[u8],
5731    x: &[f32],
5732    rows: usize,
5733    cols: usize,
5734    out: &mut [f32],
5735    pool: Option<&Pool>,
5736) {
5737    const TILE: usize = cortiq_core::quant::Q1T_TILE;
5738    let gpr = cols / GROUP_SIZE;
5739    let (rp_off, ent_off, has_ov) = q1t_overlay(bytes, rows * gpr * TILE, rows);
5740    if !has_ov {
5741        return;
5742    }
5743    let out_addr = SendMut(out.as_mut_ptr());
5744    let run = move |start: usize, end: usize| {
5745        for r in start..end {
5746            let corr = q1t_row_outlier_correction(bytes, r, rp_off, ent_off, has_ov, x);
5747            // SAFETY: disjoint rows; add onto the base the GPU already wrote.
5748            unsafe { *out_addr.at(r) += corr };
5749        }
5750    };
5751    dispatch_rows(pool, rows, &run);
5752}
5753
5754/// Q1T row range via the A8W8 int8 path — shared activation split,
5755/// per-row: base SDOT dot + outlier correction + overlay.
5756#[allow(clippy::too_many_arguments)]
5757fn q1t_range_a8w8(
5758    bytes: &[u8],
5759    gpr: usize,
5760    rp_off: usize,
5761    ent_off: usize,
5762    has_ov: bool,
5763    act: &SplitAct,
5764    x: &[f32],
5765    out: SendMut,
5766    start: usize,
5767    end: usize,
5768) {
5769    for r in start..end {
5770        let mut acc = q1t_dot_row_i8(bytes, r, gpr, &act.xq) * act.sx;
5771        for &(j, xv) in &act.outliers {
5772            acc += q1t_base_weight(bytes, r, gpr, j) * xv;
5773        }
5774        acc += q1t_row_outlier_correction(bytes, r, rp_off, ent_off, has_ov, x);
5775        // SAFETY: disjoint row ranges per worker.
5776        unsafe { *out.at(r) = acc };
5777    }
5778}
5779
5780/// Q1T row range via the f32 path (no SDOT) — for matvec_many batched
5781/// dispatch when a8w8 is unavailable.
5782#[allow(clippy::too_many_arguments)]
5783fn q1t_range_f32_batch(
5784    bytes: &[u8],
5785    gpr: usize,
5786    rp_off: usize,
5787    ent_off: usize,
5788    has_ov: bool,
5789    x: &[f32],
5790    out: SendMut,
5791    start: usize,
5792    end: usize,
5793) {
5794    const TILE: usize = cortiq_core::quant::Q1T_TILE;
5795    let mut sg = [0f32; GROUP_SIZE];
5796    for r in start..end {
5797        let mut acc = 0f32;
5798        for g in 0..gpr {
5799            let off = (r * gpr + g) * TILE;
5800            let s = f16_to_f32(u16::from_le_bytes([bytes[off], bytes[off + 1]]));
5801            let codes = &bytes[off + 2..off + TILE];
5802            let xg = &x[g * GROUP_SIZE..g * GROUP_SIZE + GROUP_SIZE];
5803            for bi in 0..6 {
5804                sg[bi * 5..bi * 5 + 5].copy_from_slice(&SIGN5[codes[bi] as usize]);
5805            }
5806            let lut = &SIGN5[codes[6] as usize];
5807            sg[30] = lut[0];
5808            sg[31] = lut[1];
5809            let mut gsum = 0f32;
5810            for k in 0..GROUP_SIZE {
5811                gsum += sg[k] * xg[k];
5812            }
5813            acc += s * gsum;
5814        }
5815        acc += q1t_row_outlier_correction(bytes, r, rp_off, ent_off, has_ov, x);
5816        // SAFETY: disjoint row ranges per worker.
5817        unsafe { *out.at(r) = acc };
5818    }
5819}
5820
5821/// Ternary (q1t) matvec — decode+dot straight from mmap, one group at a time:
5822/// no per-ROW buffer, no division (the sign LUT), and a tiny per-group sign
5823/// buffer so the 32-wide dot vectorizes. This is the decode hot path.
5824fn q1t_matvec(
5825    bytes: &[u8],
5826    x: &[f32],
5827    rows: usize,
5828    cols: usize,
5829    out: &mut [f32],
5830    pool: Option<&Pool>,
5831) {
5832    debug_assert_eq!(out.len(), rows);
5833    const TILE: usize = cortiq_core::quant::Q1T_TILE;
5834    let gpr = cols / GROUP_SIZE;
5835    let (rp_off, ent_off, has_ov) = q1t_overlay(bytes, rows * gpr * TILE, rows);
5836    let out_addr = SendMut(out.as_mut_ptr());
5837    // int8 SDOT base dot (ARM dotprod): ~4× the f32 arithmetic. x → i8 once
5838    // (`split_act`), activation outliers added back exactly in f32, weight
5839    // overlay on top. ARM SDOT / x86 AVX2; CMF_SDOT=0 keeps the exact f32 path.
5840    if a8w8_enabled() {
5841        let act = split_act(x);
5842        let act = &act;
5843        let run = move |start: usize, end: usize| {
5844            for r in start..end {
5845                let mut acc = q1t_dot_row_i8(bytes, r, gpr, &act.xq) * act.sx;
5846                for &(j, xv) in &act.outliers {
5847                    acc += q1t_base_weight(bytes, r, gpr, j) * xv;
5848                }
5849                acc += q1t_row_outlier_correction(bytes, r, rp_off, ent_off, has_ov, x);
5850                // SAFETY: disjoint row ranges per worker.
5851                unsafe { *out_addr.at(r) = acc };
5852            }
5853        };
5854        dispatch_rows(pool, rows, &run);
5855        return;
5856    }
5857    let run = move |start: usize, end: usize| {
5858        // Per-group signs, unpacked contiguously so the dot below is a clean
5859        // 32-wide reduction the autovectorizer turns into f32x4 FMAs — the
5860        // 5-values-per-byte base-3 layout won't SIMD in place.
5861        let mut sg = [0f32; GROUP_SIZE];
5862        for r in start..end {
5863            let mut acc = 0f32;
5864            for g in 0..gpr {
5865                let off = (r * gpr + g) * TILE;
5866                let s = f16_to_f32(u16::from_le_bytes([bytes[off], bytes[off + 1]]));
5867                let codes = &bytes[off + 2..off + TILE];
5868                let xg = &x[g * GROUP_SIZE..g * GROUP_SIZE + GROUP_SIZE];
5869                for bi in 0..6 {
5870                    sg[bi * 5..bi * 5 + 5].copy_from_slice(&SIGN5[codes[bi] as usize]);
5871                }
5872                let lut = &SIGN5[codes[6] as usize];
5873                sg[30] = lut[0];
5874                sg[31] = lut[1];
5875                let mut gsum = 0f32;
5876                for k in 0..GROUP_SIZE {
5877                    gsum += sg[k] * xg[k];
5878                }
5879                acc += s * gsum;
5880            }
5881            acc += q1t_row_outlier_correction(bytes, r, rp_off, ent_off, has_ov, x);
5882            unsafe { *out_addr.at(r) = acc };
5883        }
5884    };
5885    dispatch_rows(pool, rows, &run);
5886}
5887
5888/// Fused-pair twin of `q1t_dot_row_sdot`: ONE register unpack of the
5889/// ternary codes serves BOTH activation streams (the unpack chain is
5890/// the dominant per-row cost — MTP verify pairs paid it twice). Per
5891/// stream the group order and f32 accumulation match the single-row
5892/// kernel exactly, so pair == 2×matvec bit-for-bit.
5893#[cfg(target_arch = "aarch64")]
5894#[target_feature(enable = "neon,dotprod")]
5895unsafe fn q1t_dot_row_sdot2(bytes: &[u8], r: usize, gpr: usize, xa: &[i8], xb: &[i8]) -> [f32; 2] {
5896    use core::arch::aarch64::*;
5897    use core::arch::asm;
5898    // SAFETY: same slice-length contracts as `q1t_dot_row_sdot`, ×2.
5899    unsafe {
5900        const TILE: usize = cortiq_core::quant::Q1T_TILE;
5901        let bytes_ptr = bytes.as_ptr();
5902        let row_off = r * gpr * TILE;
5903        let xp = [xa.as_ptr(), xb.as_ptr()];
5904        let mut acc = [0f32; 2];
5905        macro_rules! sdot2 {
5906            ($w0:expr, $w1:expr, $x:expr) => {{
5907                let x0 = vld1q_s8($x);
5908                let x1 = vld1q_s8($x.add(16));
5909                let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
5910                asm!(
5911                    "sdot {a0:v}.4s, {w0:v}.16b, {x0:v}.16b",
5912                    "sdot {a1:v}.4s, {w1:v}.16b, {x1:v}.16b",
5913                    a0 = inout(vreg) a0, a1 = inout(vreg) a1,
5914                    w0 = in(vreg) $w0, x0 = in(vreg) x0, w1 = in(vreg) $w1, x1 = in(vreg) x1,
5915                    options(pure, nomem, nostack),
5916                );
5917                vaddvq_s32(vaddq_s32(a0, a1))
5918            }};
5919        }
5920        let gpr2 = gpr & !1;
5921        let mut gi = 0;
5922        while gi < gpr2 {
5923            let off0 = row_off + gi * TILE;
5924            let off1 = off0 + TILE;
5925            let s0 = f16_to_f32(u16::from_le_bytes([
5926                *bytes_ptr.add(off0),
5927                *bytes_ptr.add(off0 + 1),
5928            ]));
5929            let s1 = f16_to_f32(u16::from_le_bytes([
5930                *bytes_ptr.add(off1),
5931                *bytes_ptr.add(off1 + 1),
5932            ]));
5933            let (u0_0, u1_0, u2_0, u3_0) = q1t_unpack_reg_u64s(bytes_ptr.add(off0 + 2));
5934            let (u0_1, u1_1, u2_1, u3_1) = q1t_unpack_reg_u64s(bytes_ptr.add(off1 + 2));
5935            let w0_0 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u0_0), vcreate_u64(u1_0)));
5936            let w1_0 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u2_0), vcreate_u64(u3_0)));
5937            let w0_1 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u0_1), vcreate_u64(u1_1)));
5938            let w1_1 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u2_1), vcreate_u64(u3_1)));
5939            for k in 0..2 {
5940                let d0 = sdot2!(w0_0, w1_0, xp[k].add(gi * GROUP_SIZE));
5941                let d1 = sdot2!(w0_1, w1_1, xp[k].add((gi + 1) * GROUP_SIZE));
5942                acc[k] += d0 as f32 * s0 + d1 as f32 * s1;
5943            }
5944            gi += 2;
5945        }
5946        if gi < gpr {
5947            let off = row_off + gi * TILE;
5948            let s = f16_to_f32(u16::from_le_bytes([
5949                *bytes_ptr.add(off),
5950                *bytes_ptr.add(off + 1),
5951            ]));
5952            let (u0, u1, u2, u3) = q1t_unpack_reg_u64s(bytes_ptr.add(off + 2));
5953            let w0 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u0), vcreate_u64(u1)));
5954            let w1 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u2), vcreate_u64(u3)));
5955            for k in 0..2 {
5956                let d = sdot2!(w0, w1, xp[k].add(gi * GROUP_SIZE));
5957                acc[k] += d as f32 * s;
5958            }
5959        }
5960        acc
5961    }
5962}
5963
5964/// Fused Q1T pair matvec: ONE pass over the rows serves both
5965/// activation streams — on ARM the ternary register unpack happens
5966/// once per tile pair (`q1t_dot_row_sdot2`); elsewhere the second dot
5967/// rides the row's L1-warm tile bytes. Per stream the math matches
5968/// `q1t_matvec` exactly.
5969fn q1t_matvec2(
5970    bytes: &[u8],
5971    x1: &[f32],
5972    x2: &[f32],
5973    rows: usize,
5974    cols: usize,
5975    o1: &mut [f32],
5976    o2: &mut [f32],
5977    pool: Option<&Pool>,
5978) {
5979    debug_assert_eq!(o1.len(), rows);
5980    debug_assert_eq!(o2.len(), rows);
5981    const TILE: usize = cortiq_core::quant::Q1T_TILE;
5982    let gpr = cols / GROUP_SIZE;
5983    let (rp_off, ent_off, has_ov) = q1t_overlay(bytes, rows * gpr * TILE, rows);
5984    let out1 = SendMut(o1.as_mut_ptr());
5985    let out2 = SendMut(o2.as_mut_ptr());
5986    if a8w8_enabled() {
5987        let a1 = split_act(x1);
5988        let a2 = split_act(x2);
5989        let (a1, a2) = (&a1, &a2);
5990        let run = move |start: usize, end: usize| {
5991            for r in start..end {
5992                #[cfg(target_arch = "aarch64")]
5993                // a8w8 on aarch64 ⇔ sdot_enabled(), so the kernel's
5994                // target features are present.
5995                let ds = unsafe { q1t_dot_row_sdot2(bytes, r, gpr, &a1.xq, &a2.xq) };
5996                #[cfg(not(target_arch = "aarch64"))]
5997                let ds = [
5998                    q1t_dot_row_i8(bytes, r, gpr, &a1.xq),
5999                    q1t_dot_row_i8(bytes, r, gpr, &a2.xq),
6000                ];
6001                let mut acc1 = ds[0] * a1.sx;
6002                for &(j, xv) in &a1.outliers {
6003                    acc1 += q1t_base_weight(bytes, r, gpr, j) * xv;
6004                }
6005                acc1 += q1t_row_outlier_correction(bytes, r, rp_off, ent_off, has_ov, x1);
6006                let mut acc2 = ds[1] * a2.sx;
6007                for &(j, xv) in &a2.outliers {
6008                    acc2 += q1t_base_weight(bytes, r, gpr, j) * xv;
6009                }
6010                acc2 += q1t_row_outlier_correction(bytes, r, rp_off, ent_off, has_ov, x2);
6011                // SAFETY: disjoint row ranges per worker.
6012                unsafe {
6013                    *out1.at(r) = acc1;
6014                    *out2.at(r) = acc2;
6015                }
6016            }
6017        };
6018        dispatch_rows(pool, rows, &run);
6019        return;
6020    }
6021    let run = move |start: usize, end: usize| {
6022        // Exact path (CMF_SDOT=0): unpack the sign LUT once per group,
6023        // dot both streams — same op order per stream as `q1t_matvec`.
6024        let mut sg = [0f32; GROUP_SIZE];
6025        for r in start..end {
6026            let mut acc1 = 0f32;
6027            let mut acc2 = 0f32;
6028            for g in 0..gpr {
6029                let off = (r * gpr + g) * TILE;
6030                let s = f16_to_f32(u16::from_le_bytes([bytes[off], bytes[off + 1]]));
6031                let codes = &bytes[off + 2..off + TILE];
6032                for bi in 0..6 {
6033                    sg[bi * 5..bi * 5 + 5].copy_from_slice(&SIGN5[codes[bi] as usize]);
6034                }
6035                let lut = &SIGN5[codes[6] as usize];
6036                sg[30] = lut[0];
6037                sg[31] = lut[1];
6038                let xg1 = &x1[g * GROUP_SIZE..g * GROUP_SIZE + GROUP_SIZE];
6039                let xg2 = &x2[g * GROUP_SIZE..g * GROUP_SIZE + GROUP_SIZE];
6040                let mut gsum1 = 0f32;
6041                for k in 0..GROUP_SIZE {
6042                    gsum1 += sg[k] * xg1[k];
6043                }
6044                acc1 += s * gsum1;
6045                let mut gsum2 = 0f32;
6046                for k in 0..GROUP_SIZE {
6047                    gsum2 += sg[k] * xg2[k];
6048                }
6049                acc2 += s * gsum2;
6050            }
6051            acc1 += q1t_row_outlier_correction(bytes, r, rp_off, ent_off, has_ov, x1);
6052            acc2 += q1t_row_outlier_correction(bytes, r, rp_off, ent_off, has_ov, x2);
6053            // SAFETY: disjoint row ranges per worker.
6054            unsafe {
6055                *out1.at(r) = acc1;
6056                *out2.at(r) = acc2;
6057            }
6058        }
6059    };
6060    dispatch_rows(pool, rows, &run);
6061}
6062
6063/// Ternary (q1t) matmat (prefill) — dequant each row once, dot the whole
6064/// batch against it (amortizes the per-row decode).
6065fn q1t_matmat(
6066    bytes: &[u8],
6067    xs: &[f32],
6068    b: usize,
6069    rows: usize,
6070    cols: usize,
6071    out: &mut [f32],
6072    pool: Option<&Pool>,
6073) {
6074    debug_assert_eq!(out.len(), b * rows);
6075    const TILE: usize = cortiq_core::quant::Q1T_TILE;
6076    let gpr = cols / GROUP_SIZE;
6077    let (rp_off, ent_off, has_ov) = q1t_overlay(bytes, rows * gpr * TILE, rows);
6078    let out_addr = SendMut(out.as_mut_ptr());
6079    // int8 prefill (ARM SDOT / x86 AVX2): quantize the B inputs once, unpack
6080    // each weight row's signs to i8 ONCE, then int8-dot against every input —
6081    // the row sign-decode amortizes over the whole batch. CMF_SDOT=0 → f32.
6082    if a8w8_enabled() {
6083        let acts: Vec<SplitAct> = (0..b)
6084            .map(|bi| split_act(&xs[bi * cols..(bi + 1) * cols]))
6085            .collect();
6086        let acts = &acts;
6087        let run = move |start: usize, end: usize| {
6088            let mut sg = vec![0i8; cols + 8]; // row signs, i8 (+8 unpack slack)
6089            let mut sc = vec![0f32; gpr]; // per-group scales
6090            let mut accs = vec![0f32; b]; // per-batch accumulators, reused per row
6091            for r in start..end {
6092                for g in 0..gpr {
6093                    let off = (r * gpr + g) * TILE;
6094                    sc[g] = f16_to_f32(u16::from_le_bytes([bytes[off], bytes[off + 1]]));
6095                    q1t_unpack_group_i8(
6096                        bytes.as_ptr().wrapping_add(off + 2),
6097                        &mut sg[g * GROUP_SIZE..],
6098                    );
6099                }
6100                for bi in 0..b {
6101                    let act = &acts[bi];
6102                    let mut isum = 0f32;
6103                    for g in 0..gpr {
6104                        let d = q1t_i8dot32(
6105                            sg.as_ptr().wrapping_add(g * GROUP_SIZE),
6106                            act.xq.as_ptr().wrapping_add(g * GROUP_SIZE),
6107                        );
6108                        isum += d as f32 * sc[g];
6109                    }
6110                    let mut acc = isum * act.sx;
6111                    for &(j, xv) in &act.outliers {
6112                        acc += q1t_base_weight(bytes, r, gpr, j) * xv;
6113                    }
6114                    accs[bi] = acc;
6115                }
6116                // Overlay ONCE per row for the whole batch: read each (col, val)
6117                // from mmap a single time (was b× — the re-read dominated prefill)
6118                // and fan it out over the batch via the cached inputs.
6119                if has_ov {
6120                    let (c0, c1) = (
6121                        q1t_rowptr(bytes, rp_off, r),
6122                        q1t_rowptr(bytes, rp_off, r + 1),
6123                    );
6124                    for p in c0..c1 {
6125                        let e = ent_off + p * 4;
6126                        let col = u16::from_le_bytes([bytes[e], bytes[e + 1]]) as usize;
6127                        let val = f16_to_f32(u16::from_le_bytes([bytes[e + 2], bytes[e + 3]]));
6128                        for bi in 0..b {
6129                            accs[bi] += val * xs[bi * cols + col];
6130                        }
6131                    }
6132                }
6133                for bi in 0..b {
6134                    unsafe { *out_addr.at(bi * rows + r) = accs[bi] };
6135                }
6136            }
6137        };
6138        dispatch_rows(pool, rows, &run);
6139        return;
6140    }
6141    let run = move |start: usize, end: usize| {
6142        let mut buf = vec![0f32; cols];
6143        for r in start..end {
6144            q1t_dequant_row(bytes, r, gpr, rp_off, ent_off, has_ov, &mut buf);
6145            for bi in 0..b {
6146                let xr = &xs[bi * cols..(bi + 1) * cols];
6147                let mut acc = 0f32;
6148                for j in 0..cols {
6149                    acc += buf[j] * xr[j];
6150                }
6151                unsafe { *out_addr.at(bi * rows + r) = acc };
6152            }
6153        }
6154    };
6155    dispatch_rows(pool, rows, &run);
6156}
6157
6158fn q1_matvec(
6159    bytes: &[u8],
6160    x: &[f32],
6161    rows: usize,
6162    cols: usize,
6163    out: &mut [f32],
6164    pool: Option<&Pool>,
6165) {
6166    debug_assert_eq!(out.len(), rows);
6167    let gpr = cols / GROUP_SIZE;
6168    let out_addr = SendMut(out.as_mut_ptr());
6169    if a8w8_enabled() {
6170        let act = split_act(x);
6171        let gsum = q1_group_sums(&act.xq, gpr);
6172        let (act, gsum) = (&act, &gsum);
6173        let run = move |start: usize, end: usize| {
6174            q1_range_a8w8(bytes, gpr, act, gsum, out_addr, start, end)
6175        };
6176        dispatch_rows(pool, rows, &run);
6177        return;
6178    }
6179    let run = move |start: usize, end: usize| q1_range_f32(bytes, gpr, x, out_addr, start, end);
6180    dispatch_rows(pool, rows, &run);
6181}
6182
6183/// Fused two-input q1 matvec (weights read once per pair).
6184#[allow(clippy::too_many_arguments)]
6185fn q1_matvec2(
6186    bytes: &[u8],
6187    x1: &[f32],
6188    x2: &[f32],
6189    rows: usize,
6190    cols: usize,
6191    o1: &mut [f32],
6192    o2: &mut [f32],
6193    pool: Option<&Pool>,
6194) {
6195    let gpr = cols / GROUP_SIZE;
6196    let p1 = SendMut(o1.as_mut_ptr());
6197    let p2 = SendMut(o2.as_mut_ptr());
6198    if a8w8_enabled() {
6199        let a1 = split_act(x1);
6200        let a2 = split_act(x2);
6201        let g1 = q1_group_sums(&a1.xq, gpr);
6202        let g2 = q1_group_sums(&a2.xq, gpr);
6203        let (a1, a2, g1, g2) = (&a1, &a2, &g1, &g2);
6204        let run = move |start: usize, end: usize| {
6205            for r in start..end {
6206                let mut v1 = dot_q1_row_i8(bytes, r, gpr, &a1.xq, g1) * a1.sx;
6207                let mut v2 = dot_q1_row_i8(bytes, r, gpr, &a2.xq, g2) * a2.sx;
6208                for &(j, xv) in &a1.outliers {
6209                    let (w, s) = q1_outlier(bytes, r, gpr, j);
6210                    v1 += w * s * xv;
6211                }
6212                for &(j, xv) in &a2.outliers {
6213                    let (w, s) = q1_outlier(bytes, r, gpr, j);
6214                    v2 += w * s * xv;
6215                }
6216                // SAFETY: disjoint row ranges per worker.
6217                unsafe {
6218                    *p1.at(r) = v1;
6219                    *p2.at(r) = v2;
6220                }
6221            }
6222        };
6223        dispatch_rows(pool, rows, &run);
6224        return;
6225    }
6226    let run = move |start: usize, end: usize| {
6227        for r in start..end {
6228            // SAFETY: disjoint row ranges per worker.
6229            unsafe {
6230                *p1.at(r) = q1_row_exact(bytes, r, gpr, x1);
6231                *p2.at(r) = q1_row_exact(bytes, r, gpr, x2);
6232            }
6233        }
6234    };
6235    dispatch_rows(pool, rows, &run);
6236}
6237
6238/// Batched q1 matmat: each row's tiles stream once per microbatch.
6239#[allow(clippy::too_many_arguments)]
6240fn q1_matmat(
6241    bytes: &[u8],
6242    xs_all: &[f32],
6243    b: usize,
6244    rows: usize,
6245    cols: usize,
6246    out: &mut [f32],
6247    pool: Option<&Pool>,
6248) {
6249    debug_assert_eq!(out.len(), b * rows);
6250    let gpr = cols / GROUP_SIZE;
6251    let out_addr = SendMut(out.as_mut_ptr());
6252    if a8w8_enabled() {
6253        let acts: Vec<(SplitAct, Vec<i32>)> = (0..b)
6254            .map(|bi| {
6255                let act = split_act(&xs_all[bi * cols..(bi + 1) * cols]);
6256                let gsum = q1_group_sums(&act.xq, gpr);
6257                (act, gsum)
6258            })
6259            .collect();
6260        let acts = &acts;
6261        #[cfg(target_arch = "x86_64")]
6262        let blocked_ok = avx2_enabled()
6263            && std::env::var("CMF_X86_BLOCKED")
6264                .map(|v| v != "0")
6265                .unwrap_or(true);
6266        #[cfg(target_arch = "aarch64")]
6267        let blocked_ok = sdot_enabled()
6268            && std::env::var("CMF_X86_BLOCKED")
6269                .map(|v| v != "0")
6270                .unwrap_or(true);
6271        let run = move |start: usize, end: usize| {
6272            for r in start..end {
6273                let mut bi = 0usize;
6274                // Blocked 1×4: the unpacked bit mask serves four
6275                // activation streams per group.
6276                #[cfg(target_arch = "aarch64")]
6277                if blocked_ok {
6278                    while bi + 4 <= acts.len() {
6279                        let xs = [
6280                            acts[bi].0.xq.as_slice(),
6281                            acts[bi + 1].0.xq.as_slice(),
6282                            acts[bi + 2].0.xq.as_slice(),
6283                            acts[bi + 3].0.xq.as_slice(),
6284                        ];
6285                        let gs = [
6286                            acts[bi].1.as_slice(),
6287                            acts[bi + 1].1.as_slice(),
6288                            acts[bi + 2].1.as_slice(),
6289                            acts[bi + 3].1.as_slice(),
6290                        ];
6291                        let d = unsafe { dot_q1_row_1x4_sdot(bytes, r, gpr, xs, gs) };
6292                        for k in 0..4 {
6293                            let (act, _) = &acts[bi + k];
6294                            let mut acc = d[k] * act.sx;
6295                            for &(j, xv) in &act.outliers {
6296                                let (w, sc) = q1_outlier(bytes, r, gpr, j);
6297                                acc += w * sc * xv;
6298                            }
6299                            // SAFETY: disjoint (bi, r) cells per worker.
6300                            unsafe { *out_addr.at((bi + k) * rows + r) = acc };
6301                        }
6302                        bi += 4;
6303                    }
6304                }
6305                #[cfg(target_arch = "x86_64")]
6306                if blocked_ok {
6307                    while bi + 4 <= acts.len() {
6308                        let xs = [
6309                            acts[bi].0.xq.as_slice(),
6310                            acts[bi + 1].0.xq.as_slice(),
6311                            acts[bi + 2].0.xq.as_slice(),
6312                            acts[bi + 3].0.xq.as_slice(),
6313                        ];
6314                        let gs = [
6315                            acts[bi].1.as_slice(),
6316                            acts[bi + 1].1.as_slice(),
6317                            acts[bi + 2].1.as_slice(),
6318                            acts[bi + 3].1.as_slice(),
6319                        ];
6320                        let d = unsafe {
6321                            if vnni_tiles_enabled() {
6322                                dot_q1_row_1x4_vnni(bytes, r, gpr, xs, gs)
6323                            } else {
6324                                dot_q1_row_1x4_avx2(bytes, r, gpr, xs, gs)
6325                            }
6326                        };
6327                        for k in 0..4 {
6328                            let (act, _) = &acts[bi + k];
6329                            let mut acc = d[k] * act.sx;
6330                            for &(j, xv) in &act.outliers {
6331                                let (w, sc) = q1_outlier(bytes, r, gpr, j);
6332                                acc += w * sc * xv;
6333                            }
6334                            // SAFETY: disjoint (bi, r) cells per worker.
6335                            unsafe { *out_addr.at((bi + k) * rows + r) = acc };
6336                        }
6337                        bi += 4;
6338                    }
6339                }
6340                while bi < acts.len() {
6341                    let (act, gsum) = &acts[bi];
6342                    let mut acc = dot_q1_row_i8(bytes, r, gpr, &act.xq, gsum) * act.sx;
6343                    for &(j, xv) in &act.outliers {
6344                        let (w, s) = q1_outlier(bytes, r, gpr, j);
6345                        acc += w * s * xv;
6346                    }
6347                    // SAFETY: disjoint (bi, r) cells per worker range.
6348                    unsafe { *out_addr.at(bi * rows + r) = acc };
6349                    bi += 1;
6350                }
6351            }
6352        };
6353        dispatch_rows(pool, rows, &run);
6354        return;
6355    }
6356    let run = move |start: usize, end: usize| {
6357        for r in start..end {
6358            for bi in 0..b {
6359                let x = &xs_all[bi * cols..(bi + 1) * cols];
6360                // SAFETY: disjoint (bi, r) cells per worker range.
6361                unsafe { *out_addr.at(bi * rows + r) = q1_row_exact(bytes, r, gpr, x) };
6362            }
6363        }
6364    };
6365    dispatch_rows(pool, rows, &run);
6366}
6367
6368/// Fused q4_block matvec straight from the mapped bytes. SDOT path when
6369/// dotprod is available (port of vmfcore `dot_q4_block_sdot`, measured
6370/// +23% on q4 decode): nibbles → centered i8, int8×int8 `sdot` per
6371/// 32-group, exact outlier correction — the same A8W8 contract as q8.
6372/// `CMF_SDOT=0` keeps the exact scalar path.
6373fn q4matvec(
6374    bytes: &[u8],
6375    x: &[f32],
6376    rows: usize,
6377    cols: usize,
6378    out: &mut [f32],
6379    pool: Option<&Pool>,
6380) {
6381    debug_assert_eq!(out.len(), rows);
6382    let (packed, scales) = q4_split(bytes, rows, cols);
6383    let gpr = cols / GROUP_SIZE;
6384    let out_addr = SendMut(out.as_mut_ptr());
6385
6386    if a8w8_enabled() {
6387        let act = split_act(x);
6388        let run = move |start: usize, end: usize| {
6389            q4_range_a8w8(packed, scales, gpr, cols, &act, out_addr, start, end)
6390        };
6391        dispatch_rows(pool, rows, &run);
6392        return;
6393    }
6394
6395    let run =
6396        move |start: usize, end: usize| q4_range_f32(packed, scales, gpr, x, out_addr, start, end);
6397    dispatch_rows(pool, rows, &run);
6398}
6399
6400/// One q4 row via the A8W8 int8 path — SDOT on ARM, AVX2 maddubs on
6401/// x86 (scalar fallback is unreachable: callers gate on a8w8_enabled).
6402#[inline]
6403#[allow(unreachable_code)]
6404/// One UNPACKED q4 row (centered i8 in `buf`) against four activation
6405/// streams: the 32-byte weight chunk and its abs() load once per group,
6406/// the per-group f16 scale decodes once — four maddubs+reduce chains
6407/// instead of four full (load, abs, dot) rounds.
6408#[cfg(target_arch = "x86_64")]
6409#[target_feature(enable = "avx2")]
6410unsafe fn dot_q4b_row_1x4_avx2(
6411    buf: &[u8],
6412    scales: &[u8],
6413    g0: usize,
6414    gpr: usize,
6415    xs: [&[i8]; 4],
6416) -> [f32; 4] {
6417    // SAFETY: callers uphold buffer contracts (buf.len() == gpr·32).
6418    unsafe {
6419        use core::arch::x86_64::*;
6420        let ones = _mm256_set1_epi16(1);
6421        let mut acc = [0f32; 4];
6422        for gi in 0..gpr {
6423            let s = f16_to_f32(u16::from_le_bytes([
6424                scales[(g0 + gi) * 2],
6425                scales[(g0 + gi) * 2 + 1],
6426            ]));
6427            let w = _mm256_loadu_si256(buf.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
6428            let aw = _mm256_abs_epi8(w);
6429            for (k, xq) in xs.iter().enumerate() {
6430                let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
6431                let p16 = _mm256_maddubs_epi16(aw, _mm256_sign_epi8(x, w));
6432                let d = _mm256_madd_epi16(p16, ones);
6433                let hi128 = _mm256_extracti128_si256::<1>(d);
6434                let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
6435                let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
6436                let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
6437                acc[k] += _mm_cvtsi128_si32(s32) as f32 * s;
6438            }
6439        }
6440        acc
6441    }
6442}
6443
6444/// VNNI twin of `dot_q4b_row_1x4_avx2` (see `dpbusd_hsum`).
6445#[cfg(target_arch = "x86_64")]
6446#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
6447unsafe fn dot_q4b_row_1x4_vnni(
6448    buf: &[u8],
6449    scales: &[u8],
6450    g0: usize,
6451    gpr: usize,
6452    xs: [&[i8]; 4],
6453) -> [f32; 4] {
6454    // SAFETY: callers uphold buffer contracts (buf.len() == gpr·32).
6455    unsafe {
6456        use core::arch::x86_64::*;
6457        let mut acc = [0f32; 4];
6458        for gi in 0..gpr {
6459            let s = f16_to_f32(u16::from_le_bytes([
6460                scales[(g0 + gi) * 2],
6461                scales[(g0 + gi) * 2 + 1],
6462            ]));
6463            let w = _mm256_loadu_si256(buf.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
6464            let aw = _mm256_abs_epi8(w);
6465            for (k, xq) in xs.iter().enumerate() {
6466                let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
6467                let d = dpbusd_hsum(aw, _mm256_sign_epi8(x, w));
6468                acc[k] += d as f32 * s;
6469            }
6470        }
6471        acc
6472    }
6473}
6474
6475/// The vbit flavor of the blocked 1×4: the per-activation A8W8 scale
6476/// folds in PER GROUP as `(d·sx)·s` — bit-matching the single-matvec
6477/// accumulation order (the q4_block flavor applies sx once at the end,
6478/// matching ITS single path; the two conventions are historical and
6479/// each blocked leg must mirror its own).
6480#[cfg(target_arch = "x86_64")]
6481#[target_feature(enable = "avx2")]
6482unsafe fn dot_q4b_row_1x4_sx_avx2(
6483    buf: &[u8],
6484    scales: &[u8],
6485    g0: usize,
6486    gpr: usize,
6487    xs: [&[i8]; 4],
6488    sxs: [f32; 4],
6489) -> [f32; 4] {
6490    // SAFETY: callers uphold buffer contracts (buf.len() == gpr·32).
6491    unsafe {
6492        use core::arch::x86_64::*;
6493        let ones = _mm256_set1_epi16(1);
6494        let mut acc = [0f32; 4];
6495        for gi in 0..gpr {
6496            let s = f16_to_f32(u16::from_le_bytes([
6497                scales[(g0 + gi) * 2],
6498                scales[(g0 + gi) * 2 + 1],
6499            ]));
6500            let w = _mm256_loadu_si256(buf.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
6501            let aw = _mm256_abs_epi8(w);
6502            for (k, xq) in xs.iter().enumerate() {
6503                let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
6504                let p16 = _mm256_maddubs_epi16(aw, _mm256_sign_epi8(x, w));
6505                let d = _mm256_madd_epi16(p16, ones);
6506                let hi128 = _mm256_extracti128_si256::<1>(d);
6507                let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
6508                let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
6509                let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
6510                acc[k] += (_mm_cvtsi128_si32(s32) as f32 * sxs[k]) * s;
6511            }
6512        }
6513        acc
6514    }
6515}
6516
6517/// VNNI twin of `dot_q4b_row_1x4_sx_avx2` (see `dpbusd_hsum`; the
6518/// per-group `(d·sx)·s` fold mirrors the vbit single path).
6519#[cfg(target_arch = "x86_64")]
6520#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
6521unsafe fn dot_q4b_row_1x4_sx_vnni(
6522    buf: &[u8],
6523    scales: &[u8],
6524    g0: usize,
6525    gpr: usize,
6526    xs: [&[i8]; 4],
6527    sxs: [f32; 4],
6528) -> [f32; 4] {
6529    // SAFETY: callers uphold buffer contracts (buf.len() == gpr·32).
6530    unsafe {
6531        use core::arch::x86_64::*;
6532        let mut acc = [0f32; 4];
6533        for gi in 0..gpr {
6534            let s = f16_to_f32(u16::from_le_bytes([
6535                scales[(g0 + gi) * 2],
6536                scales[(g0 + gi) * 2 + 1],
6537            ]));
6538            let w = _mm256_loadu_si256(buf.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
6539            let aw = _mm256_abs_epi8(w);
6540            for (k, xq) in xs.iter().enumerate() {
6541                let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
6542                let d = dpbusd_hsum(aw, _mm256_sign_epi8(x, w));
6543                acc[k] += (d as f32 * sxs[k]) * s;
6544            }
6545        }
6546        acc
6547    }
6548}
6549
6550#[allow(unreachable_code)]
6551fn dot_q4_row_i8(packed: &[u8], scales: &[u8], g0: usize, gpr: usize, xq: &[i8]) -> f32 {
6552    #[cfg(target_arch = "aarch64")]
6553    unsafe {
6554        return dot_q4_row_sdot(packed, scales, g0, gpr, xq);
6555    }
6556    #[cfg(target_arch = "x86_64")]
6557    unsafe {
6558        return dot_q4_row_avx2(packed, scales, g0, gpr, xq);
6559    }
6560    let mut acc = 0f32;
6561    for gi in 0..gpr {
6562        let g = g0 + gi;
6563        let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
6564        let mut d = 0i32;
6565        for (k, &b) in packed[g * 16..(g + 1) * 16].iter().enumerate() {
6566            d += ((b & 0x0F) as i32 - 8) * xq[gi * GROUP_SIZE + k * 2] as i32
6567                + (((b >> 4) & 0x0F) as i32 - 8) * xq[gi * GROUP_SIZE + k * 2 + 1] as i32;
6568        }
6569        acc += d as f32 * s;
6570    }
6571    acc
6572}
6573
6574/// Two-activation q4 row via the A8W8 int8 path (see `dot_q4_row_i8`).
6575#[inline]
6576#[allow(unreachable_code)]
6577fn dot_q4_row_i8_2(
6578    packed: &[u8],
6579    scales: &[u8],
6580    g0: usize,
6581    gpr: usize,
6582    xq1: &[i8],
6583    xq2: &[i8],
6584) -> (f32, f32) {
6585    #[cfg(target_arch = "aarch64")]
6586    unsafe {
6587        return dot_q4_row_sdot2(packed, scales, g0, gpr, xq1, xq2);
6588    }
6589    #[cfg(target_arch = "x86_64")]
6590    unsafe {
6591        return dot_q4_row_avx2_2(packed, scales, g0, gpr, xq1, xq2);
6592    }
6593    (
6594        dot_q4_row_i8(packed, scales, g0, gpr, xq1),
6595        dot_q4_row_i8(packed, scales, g0, gpr, xq2),
6596    )
6597}
6598
6599/// One q4 row range via SDOT (kernel body of `q4matvec`, extracted so
6600/// multi-matrix jobs can drive it for several tensors in one dispatch).
6601#[allow(clippy::too_many_arguments)]
6602fn q4_range_a8w8(
6603    packed: &[u8],
6604    scales: &[u8],
6605    gpr: usize,
6606    cols: usize,
6607    act: &SplitAct,
6608    out: SendMut,
6609    start: usize,
6610    end: usize,
6611) {
6612    for r in start..end {
6613        let mut acc = dot_q4_row_i8(packed, scales, r * gpr, gpr, &act.xq) * act.sx;
6614        // xq is zeroed at outlier slots — add the exact terms.
6615        for &(j, xv) in &act.outliers {
6616            let flat = r * cols + j;
6617            let byte = packed[flat / 2];
6618            let nib = if flat & 1 == 0 {
6619                byte & 0x0F
6620            } else {
6621                byte >> 4
6622            };
6623            let s = f16_to_f32(u16::from_le_bytes([
6624                scales[(flat / GROUP_SIZE) * 2],
6625                scales[(flat / GROUP_SIZE) * 2 + 1],
6626            ]));
6627            acc += ((nib as i32 - 8) as f32) * s * xv;
6628        }
6629        // SAFETY: disjoint row ranges per worker.
6630        unsafe { *out.at(r) = acc };
6631    }
6632}
6633
6634/// Two-input q4 row range via the A8W8 int8 path — kernel body of
6635/// `q4matvec2`, extracted for pair multi-matrix jobs.
6636#[allow(clippy::too_many_arguments)]
6637fn q4_range2_a8w8(
6638    packed: &[u8],
6639    scales: &[u8],
6640    gpr: usize,
6641    cols: usize,
6642    a1: &SplitAct,
6643    a2: &SplitAct,
6644    p1: SendMut,
6645    p2: SendMut,
6646    start: usize,
6647    end: usize,
6648) {
6649    for r in start..end {
6650        let (s1, s2) = dot_q4_row_i8_2(packed, scales, r * gpr, gpr, &a1.xq, &a2.xq);
6651        let mut acc1 = s1 * a1.sx;
6652        let mut acc2 = s2 * a2.sx;
6653        // xq is zeroed at outlier slots — add the exact terms.
6654        let fix = |outliers: &[(usize, f32)], acc: &mut f32| {
6655            for &(j, xv) in outliers {
6656                let flat = r * cols + j;
6657                let byte = packed[flat / 2];
6658                let nib = if flat & 1 == 0 {
6659                    byte & 0x0F
6660                } else {
6661                    byte >> 4
6662                };
6663                let s = f16_to_f32(u16::from_le_bytes([
6664                    scales[(flat / GROUP_SIZE) * 2],
6665                    scales[(flat / GROUP_SIZE) * 2 + 1],
6666                ]));
6667                *acc += ((nib as i32 - 8) as f32) * s * xv;
6668            }
6669        };
6670        fix(&a1.outliers, &mut acc1);
6671        fix(&a2.outliers, &mut acc2);
6672        // SAFETY: disjoint row ranges per worker.
6673        unsafe {
6674            *p1.at(r) = acc1;
6675            *p2.at(r) = acc2;
6676        }
6677    }
6678}
6679
6680/// Exact scalar q4 row range (same extraction, non-SDOT path).
6681fn q4_range_f32(
6682    packed: &[u8],
6683    scales: &[u8],
6684    gpr: usize,
6685    x: &[f32],
6686    out: SendMut,
6687    start: usize,
6688    end: usize,
6689) {
6690    for r in start..end {
6691        let mut acc = 0f32;
6692        for gi in 0..gpr {
6693            let g = r * gpr + gi;
6694            let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
6695            let pk = &packed[g * 16..(g + 1) * 16];
6696            let xg = &x[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
6697            let mut ga = 0f32;
6698            for (k, &b) in pk.iter().enumerate() {
6699                ga += ((b & 0x0F) as f32 - 8.0) * xg[k * 2]
6700                    + (((b >> 4) & 0x0F) as f32 - 8.0) * xg[k * 2 + 1];
6701            }
6702            acc += ga * s;
6703        }
6704        // SAFETY: disjoint row ranges per worker.
6705        unsafe { *out.at(r) = acc };
6706    }
6707}
6708
6709/// Fused two-input q4 matvec: nibbles are unpacked ONCE per group and
6710/// dotted against both activations (was: two full matvecs — double
6711/// weight traffic). Per-lane math matches `q4matvec` exactly.
6712#[allow(clippy::too_many_arguments)]
6713fn q4matvec2(
6714    bytes: &[u8],
6715    x1: &[f32],
6716    x2: &[f32],
6717    rows: usize,
6718    cols: usize,
6719    o1: &mut [f32],
6720    o2: &mut [f32],
6721    pool: Option<&Pool>,
6722) {
6723    debug_assert_eq!(o1.len(), rows);
6724    debug_assert_eq!(o2.len(), rows);
6725    let (packed, scales) = q4_split(bytes, rows, cols);
6726    let gpr = cols / GROUP_SIZE;
6727
6728    if a8w8_enabled() {
6729        let a1 = split_act(x1);
6730        let a2 = split_act(x2);
6731        let p1 = SendMut(o1.as_mut_ptr());
6732        let p2 = SendMut(o2.as_mut_ptr());
6733        let run = move |start: usize, end: usize| {
6734            q4_range2_a8w8(packed, scales, gpr, cols, &a1, &a2, p1, p2, start, end)
6735        };
6736        dispatch_rows(pool, rows, &run);
6737        return;
6738    }
6739
6740    let p1 = SendMut(o1.as_mut_ptr());
6741    let p2 = SendMut(o2.as_mut_ptr());
6742    let run = move |start: usize, end: usize| {
6743        q4_range2_f32(packed, scales, gpr, x1, x2, p1, p2, start, end)
6744    };
6745    dispatch_rows(pool, rows, &run);
6746}
6747
6748/// Two-input exact scalar q4 row range (same extraction).
6749#[allow(clippy::too_many_arguments)]
6750fn q4_range2_f32(
6751    packed: &[u8],
6752    scales: &[u8],
6753    gpr: usize,
6754    x1: &[f32],
6755    x2: &[f32],
6756    p1: SendMut,
6757    p2: SendMut,
6758    start: usize,
6759    end: usize,
6760) {
6761    for r in start..end {
6762        let (mut acc1, mut acc2) = (0f32, 0f32);
6763        for gi in 0..gpr {
6764            let g = r * gpr + gi;
6765            let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
6766            let pk = &packed[g * 16..(g + 1) * 16];
6767            let x1g = &x1[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
6768            let x2g = &x2[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
6769            let (mut g1, mut g2) = (0f32, 0f32);
6770            for (k, &b) in pk.iter().enumerate() {
6771                let wl = (b & 0x0F) as f32 - 8.0;
6772                let wh = ((b >> 4) & 0x0F) as f32 - 8.0;
6773                g1 += wl * x1g[k * 2] + wh * x1g[k * 2 + 1];
6774                g2 += wl * x2g[k * 2] + wh * x2g[k * 2 + 1];
6775            }
6776            acc1 += g1 * s;
6777            acc2 += g2 * s;
6778        }
6779        // SAFETY: disjoint row ranges per worker.
6780        unsafe {
6781            *p1.at(r) = acc1;
6782            *p2.at(r) = acc2;
6783        }
6784    }
6785}
6786
6787thread_local! {
6788    /// Per-worker decoded-row scratch for the batched q4/vbit kernels
6789    /// (centered i8 for SDOT, f32 for the exact/scalar paths).
6790    static ROW_I8: std::cell::RefCell<Vec<u8>> = const { std::cell::RefCell::new(Vec::new()) };
6791    static ROW_F32: std::cell::RefCell<Vec<f32>> = const { std::cell::RefCell::new(Vec::new()) };
6792}
6793
6794/// Batched q4 matmat: each weight row is unpacked from the mmap ONCE
6795/// and dotted against ALL b activations (prefill used to fall back to b
6796/// full matvecs — b× weight traffic and b× nibble decode). Per-position
6797/// math matches `q4matvec` exactly: same group order, same accumulation.
6798/// `out` is row-major [b, rows] like `qmatmat`.
6799#[allow(clippy::too_many_arguments)]
6800fn q4matmat(
6801    bytes: &[u8],
6802    xs_all: &[f32],
6803    b: usize,
6804    rows: usize,
6805    cols: usize,
6806    out: &mut [f32],
6807    pool: Option<&Pool>,
6808) {
6809    debug_assert_eq!(xs_all.len(), b * cols);
6810    debug_assert_eq!(out.len(), b * rows);
6811    let (packed, scales) = q4_split(bytes, rows, cols);
6812    let gpr = cols / GROUP_SIZE;
6813    let gscale = |g: usize| f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
6814
6815    if a8w8_enabled() {
6816        let acts: Vec<SplitAct> = (0..b)
6817            .map(|bi| split_act(&xs_all[bi * cols..(bi + 1) * cols]))
6818            .collect();
6819        let acts = &acts;
6820        let out_addr = SendMut(out.as_mut_ptr());
6821        let run = move |start: usize, end: usize| {
6822            ROW_I8.with(|rb| {
6823                let mut buf = rb.borrow_mut();
6824                buf.resize(cols, 0);
6825                for r in start..end {
6826                    // Unpack the row's nibbles to centered i8 once
6827                    // (element 2k = low nibble, 2k+1 = high — flat order,
6828                    // same as dot_q4_row_sdot's zip).
6829                    for gi in 0..gpr {
6830                        let g = r * gpr + gi;
6831                        for (k, &bt) in packed[g * 16..(g + 1) * 16].iter().enumerate() {
6832                            buf[gi * GROUP_SIZE + k * 2] = ((bt & 0x0F) as i32 - 8) as i8 as u8;
6833                            buf[gi * GROUP_SIZE + k * 2 + 1] =
6834                                (((bt >> 4) & 0x0F) as i32 - 8) as i8 as u8;
6835                        }
6836                    }
6837                    let mut bi = 0usize;
6838                    #[cfg(target_arch = "x86_64")]
6839                    if avx2_enabled()
6840                        && std::env::var("CMF_X86_BLOCKED")
6841                            .map(|v| v != "0")
6842                            .unwrap_or(true)
6843                    {
6844                        while bi + 4 <= acts.len() {
6845                            let xs = [
6846                                acts[bi].xq.as_slice(),
6847                                acts[bi + 1].xq.as_slice(),
6848                                acts[bi + 2].xq.as_slice(),
6849                                acts[bi + 3].xq.as_slice(),
6850                            ];
6851                            let d = unsafe {
6852                                if vnni_tiles_enabled() {
6853                                    dot_q4b_row_1x4_vnni(&buf, scales, r * gpr, gpr, xs)
6854                                } else {
6855                                    dot_q4b_row_1x4_avx2(&buf, scales, r * gpr, gpr, xs)
6856                                }
6857                            };
6858                            for k in 0..4 {
6859                                let act = &acts[bi + k];
6860                                let mut acc = d[k] * act.sx;
6861                                for &(j, xv) in &act.outliers {
6862                                    acc += (buf[j] as i8) as f32
6863                                        * gscale((r * cols + j) / GROUP_SIZE)
6864                                        * xv;
6865                                }
6866                                // SAFETY: disjoint (bi, r) cells per worker.
6867                                unsafe { *out_addr.at((bi + k) * rows + r) = acc };
6868                            }
6869                            bi += 4;
6870                        }
6871                    }
6872                    while bi < acts.len() {
6873                        let act = &acts[bi];
6874                        let mut acc = 0f32;
6875                        for gi in 0..gpr {
6876                            let d = dot_i8_i8(
6877                                &buf[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE],
6878                                &act.xq[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE],
6879                            );
6880                            acc += d as f32 * gscale(r * gpr + gi);
6881                        }
6882                        acc *= act.sx;
6883                        // xq is zeroed at outlier slots — exact terms.
6884                        for &(j, xv) in &act.outliers {
6885                            acc += (buf[j] as i8) as f32 * gscale((r * cols + j) / GROUP_SIZE) * xv;
6886                        }
6887                        // SAFETY: disjoint (bi, r) cells per worker row range.
6888                        unsafe { *out_addr.at(bi * rows + r) = acc };
6889                        bi += 1;
6890                    }
6891                }
6892            })
6893        };
6894        dispatch_rows(pool, rows, &run);
6895        return;
6896    }
6897
6898    let out_addr = SendMut(out.as_mut_ptr());
6899    let run = move |start: usize, end: usize| {
6900        ROW_F32.with(|rb| {
6901            let mut buf = rb.borrow_mut();
6902            buf.resize(cols, 0.0);
6903            for r in start..end {
6904                // Decode raw (nib − 8) values once; scales stay per-group
6905                // so the accumulation order matches q4matvec bit-for-bit.
6906                for gi in 0..gpr {
6907                    let g = r * gpr + gi;
6908                    for (k, &bt) in packed[g * 16..(g + 1) * 16].iter().enumerate() {
6909                        buf[gi * GROUP_SIZE + k * 2] = (bt & 0x0F) as f32 - 8.0;
6910                        buf[gi * GROUP_SIZE + k * 2 + 1] = ((bt >> 4) & 0x0F) as f32 - 8.0;
6911                    }
6912                }
6913                for bi in 0..b {
6914                    let x = &xs_all[bi * cols..(bi + 1) * cols];
6915                    let mut acc = 0f32;
6916                    for gi in 0..gpr {
6917                        let mut ga = 0f32;
6918                        // Pairwise (lo + hi) addition, matching
6919                        // q4matvec's `ga += lo·x + hi·x` shape exactly —
6920                        // a flat one-per-element loop rounds differently
6921                        // and broke bit-parity on the scalar (x86) path.
6922                        for k in 0..GROUP_SIZE / 2 {
6923                            let e = gi * GROUP_SIZE + k * 2;
6924                            ga += buf[e] * x[e] + buf[e + 1] * x[e + 1];
6925                        }
6926                        acc += ga * gscale(r * gpr + gi);
6927                    }
6928                    // SAFETY: disjoint (bi, r) cells per worker row range.
6929                    unsafe { *out_addr.at(bi * rows + r) = acc };
6930                }
6931            }
6932        })
6933    };
6934    dispatch_rows(pool, rows, &run);
6935}
6936
6937/// Batched vbit matmat: each variable-bit row is decoded from the mmap
6938/// ONCE for the whole microbatch. Same per-position math as
6939/// `vbitmatvec` (SDOT A8W8 with exact outliers / exact f32 for b=8 rows
6940/// and the scalar path).
6941#[allow(clippy::too_many_arguments)]
6942fn vbitmatmat(
6943    bytes: &[u8],
6944    offsets: &[usize],
6945    xs_all: &[f32],
6946    b: usize,
6947    rows: usize,
6948    cols: usize,
6949    out: &mut [f32],
6950    pool: Option<&Pool>,
6951) {
6952    debug_assert_eq!(xs_all.len(), b * cols);
6953    debug_assert_eq!(out.len(), b * rows);
6954    debug_assert_eq!(offsets.len(), rows + 1);
6955    let ng = cols / GROUP_SIZE;
6956    let bits = &bytes[..rows];
6957    let sc_off = rows;
6958    let gscale = |r: usize, g: usize| {
6959        let so = (r * ng + g) * 2;
6960        f16_to_f32(u16::from_le_bytes([
6961            bytes[sc_off + so],
6962            bytes[sc_off + so + 1],
6963        ]))
6964    };
6965
6966    // Decode row r's raw (u − L) values into `dst` (f32, unscaled).
6967    let decode_f32 = |r: usize, dst: &mut [f32]| {
6968        let bw = bits[r] as usize;
6969        let l = ((1i32 << (bw - 1)) - 1) as f32;
6970        let data = &bytes[offsets[r]..offsets[r + 1]];
6971        let (mut acc, mut nbits, mut idx) = (0u64, 0usize, 0usize);
6972        for d in dst.iter_mut() {
6973            while nbits < bw {
6974                acc = (acc << 8) | data[idx] as u64;
6975                idx += 1;
6976                nbits += 8;
6977            }
6978            let u = ((acc >> (nbits - bw)) & ((1u64 << bw) - 1)) as f32;
6979            nbits -= bw;
6980            *d = u - l;
6981        }
6982    };
6983
6984    if a8w8_enabled() {
6985        let acts: Vec<SplitAct> = (0..b)
6986            .map(|bi| split_act(&xs_all[bi * cols..(bi + 1) * cols]))
6987            .collect();
6988        let acts = &acts;
6989        let out_addr = SendMut(out.as_mut_ptr());
6990        let run = move |start: usize, end: usize| {
6991            for r in start..end {
6992                let bw = bits[r] as usize;
6993                if bw == 8 {
6994                    // u−L reaches 128 → no i8 path; decode once, exact
6995                    // f32 dots for every position (same as vbitmatvec).
6996                    ROW_F32.with(|rb| {
6997                        let mut buf = rb.borrow_mut();
6998                        buf.resize(cols, 0.0);
6999                        decode_f32(r, &mut buf);
7000                        for bi in 0..b {
7001                            let x = &xs_all[bi * cols..(bi + 1) * cols];
7002                            let mut dot = 0f32;
7003                            for g in 0..ng {
7004                                let mut gd = 0f32;
7005                                for k in 0..GROUP_SIZE {
7006                                    gd += buf[g * GROUP_SIZE + k] * x[g * GROUP_SIZE + k];
7007                                }
7008                                dot += gd * gscale(r, g);
7009                            }
7010                            // SAFETY: disjoint (bi, r) cells per worker range.
7011                            unsafe { *out_addr.at(bi * rows + r) = dot };
7012                        }
7013                    });
7014                    continue;
7015                }
7016                let l = (1i32 << (bw - 1)) - 1;
7017                let data = &bytes[offsets[r]..offsets[r + 1]];
7018                ROW_I8.with(|rb| {
7019                    let mut buf = rb.borrow_mut();
7020                    buf.resize(cols, 0);
7021                    #[inline(always)]
7022                    fn fill<const B: usize>(data: &[u8], l: i32, buf: &mut [u8]) {
7023                        for (blk, chunk) in buf.chunks_exact_mut(8).enumerate() {
7024                            let u = unpack8::<B>(&data[blk * B..]);
7025                            for k in 0..8 {
7026                                chunk[k] = (u[k] - l) as i8 as u8;
7027                            }
7028                        }
7029                    }
7030                    match bw {
7031                        3 => fill::<3>(data, l, &mut buf),
7032                        4 => vbit_fill4(data, &mut buf),
7033                        5 => fill::<5>(data, l, &mut buf),
7034                        6 => fill::<6>(data, l, &mut buf),
7035                        _ => unreachable!("vbit bit-width {bw} (validated at load)"),
7036                    }
7037                    let mut bi = 0usize;
7038                    // The vbit scale table shares q4_block's layout
7039                    // (contiguous f16 per (row·ng + g)), so the same
7040                    // blocked 1×4 kernel serves the decoded row.
7041                    #[cfg(target_arch = "x86_64")]
7042                    if avx2_enabled()
7043                        && std::env::var("CMF_X86_BLOCKED")
7044                            .map(|v| v != "0")
7045                            .unwrap_or(true)
7046                    {
7047                        while bi + 4 <= acts.len() {
7048                            let xs = [
7049                                acts[bi].xq.as_slice(),
7050                                acts[bi + 1].xq.as_slice(),
7051                                acts[bi + 2].xq.as_slice(),
7052                                acts[bi + 3].xq.as_slice(),
7053                            ];
7054                            let sxs = [
7055                                acts[bi].sx,
7056                                acts[bi + 1].sx,
7057                                acts[bi + 2].sx,
7058                                acts[bi + 3].sx,
7059                            ];
7060                            let d = unsafe {
7061                                if vnni_tiles_enabled() {
7062                                    dot_q4b_row_1x4_sx_vnni(
7063                                        &buf,
7064                                        &bytes[sc_off..],
7065                                        r * ng,
7066                                        ng,
7067                                        xs,
7068                                        sxs,
7069                                    )
7070                                } else {
7071                                    dot_q4b_row_1x4_sx_avx2(
7072                                        &buf,
7073                                        &bytes[sc_off..],
7074                                        r * ng,
7075                                        ng,
7076                                        xs,
7077                                        sxs,
7078                                    )
7079                                }
7080                            };
7081                            for k in 0..4 {
7082                                let act = &acts[bi + k];
7083                                let mut dot = d[k];
7084                                for &(j, xv) in &act.outliers {
7085                                    dot += (buf[j] as i8) as f32 * gscale(r, j / GROUP_SIZE) * xv;
7086                                }
7087                                // SAFETY: disjoint (bi, r) cells per worker.
7088                                unsafe { *out_addr.at((bi + k) * rows + r) = dot };
7089                            }
7090                            bi += 4;
7091                        }
7092                    }
7093                    while bi < acts.len() {
7094                        let act = &acts[bi];
7095                        let mut dot = 0f32;
7096                        for g in 0..ng {
7097                            let d = dot_i8_i8(
7098                                &buf[g * GROUP_SIZE..(g + 1) * GROUP_SIZE],
7099                                &act.xq[g * GROUP_SIZE..(g + 1) * GROUP_SIZE],
7100                            ) as f32
7101                                * act.sx;
7102                            dot += d * gscale(r, g);
7103                        }
7104                        for &(j, xv) in &act.outliers {
7105                            dot += (buf[j] as i8) as f32 * gscale(r, j / GROUP_SIZE) * xv;
7106                        }
7107                        // SAFETY: disjoint (bi, r) cells per worker range.
7108                        unsafe { *out_addr.at(bi * rows + r) = dot };
7109                        bi += 1;
7110                    }
7111                });
7112            }
7113        };
7114        dispatch_rows(pool, rows, &run);
7115        return;
7116    }
7117
7118    let out_addr = SendMut(out.as_mut_ptr());
7119    let run = move |start: usize, end: usize| {
7120        ROW_F32.with(|rb| {
7121            let mut buf = rb.borrow_mut();
7122            buf.resize(cols, 0.0);
7123            for r in start..end {
7124                decode_f32(r, &mut buf);
7125                for bi in 0..b {
7126                    let x = &xs_all[bi * cols..(bi + 1) * cols];
7127                    let mut dot = 0f32;
7128                    for g in 0..ng {
7129                        let mut gd = 0f32;
7130                        for k in 0..GROUP_SIZE {
7131                            gd += buf[g * GROUP_SIZE + k] * x[g * GROUP_SIZE + k];
7132                        }
7133                        dot += gd * gscale(r, g);
7134                    }
7135                    // SAFETY: disjoint (bi, r) cells per worker range.
7136                    unsafe { *out_addr.at(bi * rows + r) = dot };
7137                }
7138            }
7139        })
7140    };
7141    dispatch_rows(pool, rows, &run);
7142}
7143
7144/// Build a GPU batch job for a q8-family mapped tensor (primary
7145/// shard): prescaled input + directory coordinates. None → not
7146/// GPU-eligible, caller stays on the CPU.
7147pub(crate) fn gpu_batch_job<'a>(
7148    t: &'a QTensor,
7149    x: &[f32],
7150) -> Option<(std::sync::Arc<CmfModel>, crate::gpu::BatchJob<'a>)> {
7151    match t {
7152        QTensor::Mapped {
7153            model,
7154            idx,
7155            dtype: dt @ (TensorDtype::Q8Row | TensorDtype::Q8_2f),
7156            rows,
7157            cols,
7158            row_scale,
7159            col_field,
7160            ..
7161        } => Some((
7162            model.clone(),
7163            crate::gpu::BatchJob {
7164                idx: *idx,
7165                rows: *rows,
7166                cols: *cols,
7167                row_scale,
7168                xs: prescale(x, col_field, *dt).into_owned(),
7169                layout: crate::gpu::BatchLayout::Q8,
7170            },
7171        )),
7172        // q1: raw f32 activations, tile-embedded scales.
7173        QTensor::Mapped {
7174            model,
7175            idx,
7176            dtype: TensorDtype::Q1,
7177            rows,
7178            cols,
7179            ..
7180        } => Some((
7181            model.clone(),
7182            crate::gpu::BatchJob {
7183                idx: *idx,
7184                rows: *rows,
7185                cols: *cols,
7186                row_scale: &[],
7187                xs: x.to_vec(),
7188                layout: crate::gpu::BatchLayout::Q1,
7189            },
7190        )),
7191        _ => None,
7192    }
7193}
7194
7195thread_local! {
7196    static PRESCALE_BUF1: std::cell::RefCell<Vec<f32>> = const { std::cell::RefCell::new(Vec::new()) };
7197    static PRESCALE_BUF2: std::cell::RefCell<Vec<f32>> = const { std::cell::RefCell::new(Vec::new()) };
7198}
7199
7200pub(crate) fn prescale<'a>(
7201    x: &'a [f32],
7202    col_field: &[f32],
7203    dtype: TensorDtype,
7204) -> std::borrow::Cow<'a, [f32]> {
7205    if dtype == TensorDtype::Q8_2f {
7206        x.iter().zip(col_field).map(|(a, c)| a * c).collect()
7207    } else {
7208        std::borrow::Cow::Borrowed(x)
7209    }
7210}
7211
7212/// θ col-field fold for q8_2f activations. Borrowed pass-through for
7213/// every other dtype, using thread-local buffers to eliminate per-matvec allocations.
7214pub(crate) fn prescale_with<R, F: FnOnce(&[f32]) -> R>(
7215    x: &[f32],
7216    col_field: &[f32],
7217    dtype: TensorDtype,
7218    buf_id: u8,
7219    f: F,
7220) -> R {
7221    if dtype == TensorDtype::Q8_2f {
7222        if buf_id == 1 {
7223            PRESCALE_BUF1.with(|b| {
7224                let mut buf = b.borrow_mut();
7225                buf.clear();
7226                buf.extend(x.iter().zip(col_field).map(|(a, c)| a * c));
7227                f(&buf)
7228            })
7229        } else {
7230            PRESCALE_BUF2.with(|b| {
7231                let mut buf = b.borrow_mut();
7232                buf.clear();
7233                buf.extend(x.iter().zip(col_field).map(|(a, c)| a * c));
7234                f(&buf)
7235            })
7236        }
7237    } else {
7238        f(x)
7239    }
7240}
7241
7242// ───────────────────── x86-64 AVX2 kernels (roadmap этап 2) ─────────────────────
7243
7244/// AVX2+FMA available? Default ON when the CPU supports both;
7245/// `CMF_AVX2=0` disables (falls back to the autovectorized loops).
7246#[cfg(target_arch = "x86_64")]
7247pub(crate) fn avx2_enabled() -> bool {
7248    use std::sync::OnceLock;
7249    static ON: OnceLock<bool> = OnceLock::new();
7250    *ON.get_or_init(|| {
7251        std::env::var("CMF_AVX2").map(|v| v != "0").unwrap_or(true)
7252            && std::arch::is_x86_feature_detected!("avx2")
7253            && std::arch::is_x86_feature_detected!("fma")
7254    })
7255}
7256
7257/// AVX2 A8W8 allowed? The quantized-activation contract is switched by
7258/// the SAME env as the ARM SDOT path: `CMF_SDOT=0` keeps exact kernels
7259/// (the golden-parity exact gate relies on it) — AVX2 f32 kernels stay
7260/// active either way, they are exact (regrouped sums only).
7261#[cfg(target_arch = "x86_64")]
7262fn avx2_a8w8_enabled() -> bool {
7263    use std::sync::OnceLock;
7264    static ON: OnceLock<bool> = OnceLock::new();
7265    *ON.get_or_init(|| {
7266        avx2_enabled() && std::env::var("CMF_SDOT").map(|v| v != "0").unwrap_or(true)
7267    })
7268}
7269
7270/// A8W8 quantized-activation path available on THIS machine? One
7271/// switch across architectures: ARM dotprod (CMF_SDOT) or x86 AVX2
7272/// (CMF_AVX2 + the same CMF_SDOT exact-contract override).
7273#[inline]
7274pub(crate) fn a8w8_enabled() -> bool {
7275    #[cfg(target_arch = "aarch64")]
7276    {
7277        sdot_enabled()
7278    }
7279    #[cfg(target_arch = "x86_64")]
7280    {
7281        avx2_a8w8_enabled()
7282    }
7283    #[cfg(not(any(target_arch = "aarch64", target_arch = "x86_64")))]
7284    {
7285        false
7286    }
7287}
7288
7289/// int8·int8 dot dispatch: SDOT on ARM; AVX-512 VNNI (vpdpbusd) or AVX2
7290/// maddubs on x86. Callers are gated by `a8w8_enabled()`.
7291#[inline]
7292#[allow(unreachable_code)]
7293fn dot_i8_i8(w: &[u8], xq: &[i8]) -> i32 {
7294    #[cfg(target_arch = "aarch64")]
7295    unsafe {
7296        return dot_i8_sdot(w, xq);
7297    }
7298    #[cfg(target_arch = "x86_64")]
7299    unsafe {
7300        if avx512vnni_enabled() {
7301            return dot_i8_i8_vnni(w, xq);
7302        }
7303        return dot_i8_i8_avx2(w, xq);
7304    }
7305    w.iter()
7306        .zip(xq)
7307        .map(|(&a, &b)| (a as i8) as i32 * b as i32)
7308        .sum()
7309}
7310
7311/// AVX-512 VNNI available? (F+BW+VL+VNNI; `CMF_AVX512=0` falls back to
7312/// AVX2.) VL matters: short 32-byte groups (q4/vbit) ride the 256-bit
7313/// `vpdpbusd` encoding.
7314#[cfg(target_arch = "x86_64")]
7315fn avx512vnni_enabled() -> bool {
7316    use std::sync::OnceLock;
7317    static ON: OnceLock<bool> = OnceLock::new();
7318    *ON.get_or_init(|| {
7319        std::env::var("CMF_AVX512")
7320            .map(|v| v != "0")
7321            .unwrap_or(true)
7322            && std::arch::is_x86_feature_detected!("avx512f")
7323            && std::arch::is_x86_feature_detected!("avx512bw")
7324            && std::arch::is_x86_feature_detected!("avx512vl")
7325            && std::arch::is_x86_feature_detected!("avx512vnni")
7326    })
7327}
7328
7329/// Grouped-codec VNNI arms (the q4t/q4b/q1/q1t tile kernels): default
7330/// ON where AVX-512 VNNI exists (`CMF_VNNI_TILES=0` opt-out). Measured
7331/// on Ryzen 7950X (Zen4, 3 alternating process pairs, blocked GEMM
7332/// 4864×896 b=256): q4t 63→68 GF/s (+8%), q1 53→56 (+6%), q4b 72→75
7333/// (+4%) — consistent, no leg regressed. The tile kernels keep a
7334/// horizontal reduce per 32-weight group, so the `vpdpbusd` saving is
7335/// smaller than the long-dot q8 win (+13%), but it is real and free.
7336#[cfg(target_arch = "x86_64")]
7337fn vnni_tiles_enabled() -> bool {
7338    use std::sync::OnceLock;
7339    static ON: OnceLock<bool> = OnceLock::new();
7340    *ON.get_or_init(|| {
7341        std::env::var("CMF_VNNI_TILES")
7342            .map(|v| v != "0")
7343            .unwrap_or(true)
7344            && avx512vnni_enabled()
7345    })
7346}
7347
7348/// One 256-bit u8×i8 dot → i32 via `vpdpbusd` into a fresh accumulator
7349/// plus the same horizontal reduce the AVX2 kernels use. Products are
7350/// bounded (|w| ≤ 8 or ≤ 1), so maddubs never saturated — the i32 sum
7351/// is bit-identical to the maddubs+madd pair it replaces.
7352#[cfg(target_arch = "x86_64")]
7353#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
7354#[inline]
7355unsafe fn dpbusd_hsum(aw: core::arch::x86_64::__m256i, xs: core::arch::x86_64::__m256i) -> i32 {
7356    // SAFETY: pure register math.
7357    unsafe {
7358        use core::arch::x86_64::*;
7359        let d = _mm256_dpbusd_epi32(_mm256_setzero_si256(), aw, xs);
7360        let hi128 = _mm256_extracti128_si256::<1>(d);
7361        let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
7362        let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
7363        let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
7364        _mm_cvtsi128_si32(s32)
7365    }
7366}
7367
7368/// int8·int8 via AVX-512 VNNI: `vpdpbusd` fuses the maddubs+madd+add
7369/// triple into one u8×i8 dot-accumulate. AVX-512 has no vpsignb, so the
7370/// |w|·sign(x,w) trick becomes |w| × (x negated where w<0) via a mask
7371/// subtract — w==0 lanes contribute 0 through |w|=0 either way.
7372#[cfg(target_arch = "x86_64")]
7373#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
7374unsafe fn dot_i8_i8_vnni(w: &[u8], xq: &[i8]) -> i32 {
7375    // SAFETY: callers uphold slice-length contracts (see call sites).
7376    unsafe {
7377        use core::arch::x86_64::*;
7378        let n = w.len();
7379        let mut j = 0usize;
7380        let mut total: i32;
7381        // 4 independent accumulators: vpdpbusd is its own loop-carried
7382        // dependency (~5-cycle latency) — a single-acc loop runs
7383        // latency-bound and LOSES to the AVX2 maddubs kernel, measured
7384        // on Granite Rapids.
7385        {
7386            #[inline(always)]
7387            unsafe fn step(
7388                w: *const u8,
7389                x: *const i8,
7390                acc: core::arch::x86_64::__m512i,
7391            ) -> core::arch::x86_64::__m512i {
7392                unsafe {
7393                    use core::arch::x86_64::*;
7394                    let wv = _mm512_loadu_si512(w as *const _);
7395                    let xv = _mm512_loadu_si512(x as *const _);
7396                    let aw = _mm512_abs_epi8(wv);
7397                    let neg = _mm512_movepi8_mask(wv);
7398                    let sx = _mm512_mask_sub_epi8(xv, neg, _mm512_setzero_si512(), xv);
7399                    _mm512_dpbusd_epi32(acc, aw, sx)
7400                }
7401            }
7402            let (mut a0, mut a1, mut a2, mut a3) = (
7403                _mm512_setzero_si512(),
7404                _mm512_setzero_si512(),
7405                _mm512_setzero_si512(),
7406                _mm512_setzero_si512(),
7407            );
7408            while j + 256 <= n {
7409                a0 = step(w.as_ptr().add(j), xq.as_ptr().add(j), a0);
7410                a1 = step(w.as_ptr().add(j + 64), xq.as_ptr().add(j + 64), a1);
7411                a2 = step(w.as_ptr().add(j + 128), xq.as_ptr().add(j + 128), a2);
7412                a3 = step(w.as_ptr().add(j + 192), xq.as_ptr().add(j + 192), a3);
7413                j += 256;
7414            }
7415            while j + 64 <= n {
7416                a0 = step(w.as_ptr().add(j), xq.as_ptr().add(j), a0);
7417                j += 64;
7418            }
7419            let s01 = _mm512_add_epi32(a0, a1);
7420            let s23 = _mm512_add_epi32(a2, a3);
7421            total = _mm512_reduce_add_epi32(_mm512_add_epi32(s01, s23));
7422        }
7423        // 32-wide (q4/vbit groups are exactly 32 bytes).
7424        if j + 32 <= n {
7425            let wv = _mm256_loadu_si256(w.as_ptr().add(j) as *const __m256i);
7426            let xv = _mm256_loadu_si256(xq.as_ptr().add(j) as *const __m256i);
7427            let d = _mm256_dpbusd_epi32(
7428                _mm256_setzero_si256(),
7429                _mm256_abs_epi8(wv),
7430                _mm256_sign_epi8(xv, wv),
7431            );
7432            let hi128 = _mm256_extracti128_si256::<1>(d);
7433            let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
7434            let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
7435            let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
7436            total += _mm_cvtsi128_si32(s32);
7437            j += 32;
7438        }
7439        while j < n {
7440            total += (w[j] as i8) as i32 * xq[j] as i32;
7441            j += 1;
7442        }
7443        total
7444    }
7445}
7446
7447/// i8 row · f32 x via AVX2/FMA (x86 mirror of `dot_i8_f32_neon`).
7448#[cfg(target_arch = "x86_64")]
7449#[target_feature(enable = "avx2,fma")]
7450unsafe fn dot_i8_f32_avx2(w: &[u8], x: &[f32]) -> f32 {
7451    // SAFETY: callers uphold slice-length contracts (see call sites).
7452    unsafe {
7453        use core::arch::x86_64::*;
7454        let n = x.len();
7455        let wp = w.as_ptr();
7456        let xp = x.as_ptr();
7457        let (mut a0, mut a1) = (_mm256_setzero_ps(), _mm256_setzero_ps());
7458        let mut j = 0usize;
7459        while j + 16 <= n {
7460            let wb = _mm_loadu_si128(wp.add(j) as *const __m128i);
7461            let lo = _mm256_cvtepi8_epi32(wb);
7462            let hi = _mm256_cvtepi8_epi32(_mm_srli_si128::<8>(wb));
7463            a0 = _mm256_fmadd_ps(_mm256_cvtepi32_ps(lo), _mm256_loadu_ps(xp.add(j)), a0);
7464            a1 = _mm256_fmadd_ps(_mm256_cvtepi32_ps(hi), _mm256_loadu_ps(xp.add(j + 8)), a1);
7465            j += 16;
7466        }
7467        let acc = _mm256_add_ps(a0, a1);
7468        let hi128 = _mm256_extractf128_ps::<1>(acc);
7469        let s128 = _mm_add_ps(_mm256_castps256_ps128(acc), hi128);
7470        let s64 = _mm_add_ps(s128, _mm_movehl_ps(s128, s128));
7471        let s32 = _mm_add_ss(s64, _mm_shuffle_ps::<1>(s64, s64));
7472        let mut sum = _mm_cvtss_f32(s32);
7473        while j < n {
7474            sum += (*wp.add(j) as i8) as f32 * *xp.add(j);
7475            j += 1;
7476        }
7477        sum
7478    }
7479}
7480
7481/// int8(weight)·int8(activation) → i32 via AVX2 maddubs — the x86
7482/// analogue of the SDOT A8W8 path. `maddubs` takes u8×i8, so the
7483/// standard sign trick applies: |w| × sign(x, w) ≡ w × x per lane.
7484/// Pair saturation is safe: |w|≤128, |x|≤127 → 2·128·127 < 32767.
7485#[cfg(target_arch = "x86_64")]
7486#[target_feature(enable = "avx2")]
7487unsafe fn dot_i8_i8_avx2(w: &[u8], xq: &[i8]) -> i32 {
7488    // SAFETY: callers uphold slice-length contracts (see call sites).
7489    unsafe {
7490        use core::arch::x86_64::*;
7491        let n = w.len();
7492        let ones = _mm256_set1_epi16(1);
7493        let mut acc = _mm256_setzero_si256();
7494        let mut j = 0usize;
7495        while j + 32 <= n {
7496            let wv = _mm256_loadu_si256(w.as_ptr().add(j) as *const __m256i);
7497            let xv = _mm256_loadu_si256(xq.as_ptr().add(j) as *const __m256i);
7498            let p16 = _mm256_maddubs_epi16(_mm256_abs_epi8(wv), _mm256_sign_epi8(xv, wv));
7499            acc = _mm256_add_epi32(acc, _mm256_madd_epi16(p16, ones));
7500            j += 32;
7501        }
7502        let hi128 = _mm256_extracti128_si256::<1>(acc);
7503        let s128 = _mm_add_epi32(_mm256_castsi256_si128(acc), hi128);
7504        let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
7505        let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
7506        let mut s = _mm_cvtsi128_si32(s32);
7507        while j < n {
7508            s += (w[j] as i8) as i32 * xq[j] as i32;
7509            j += 1;
7510        }
7511        s
7512    }
7513}
7514
7515/// smmla 2×4: one instruction covers a 2-row × 2-activation × 8-deep
7516/// tile (32 MACs vs sdot's 16) — the weight pair loads once per 8-k
7517/// slice as a combined 2×8 register and meets two activation pairs.
7518#[cfg(target_arch = "aarch64")]
7519#[target_feature(enable = "neon,i8mm")]
7520unsafe fn dot_i8_smmla_2x4(w0: &[u8], w1: &[u8], xs: [&[i8]; 4]) -> [[i32; 4]; 2] {
7521    // SAFETY: callers uphold slice-length contracts.
7522    unsafe {
7523        use core::arch::aarch64::*;
7524        use core::arch::asm;
7525        let n = w0.len();
7526        let w0p = w0.as_ptr() as *const i8;
7527        let w1p = w1.as_ptr() as *const i8;
7528        // acc01 holds [c(r0,x0) c(r0,x1) c(r1,x0) c(r1,x1)]; acc23 the
7529        // same for x2/x3.
7530        let mut acc01 = vdupq_n_s32(0);
7531        let mut acc23 = vdupq_n_s32(0);
7532        let mut i = 0usize;
7533        while i + 8 <= n {
7534            let wa = vcombine_s8(vld1_s8(w0p.add(i)), vld1_s8(w1p.add(i)));
7535            let xb01 = vcombine_s8(
7536                vld1_s8(xs[0].as_ptr().add(i)),
7537                vld1_s8(xs[1].as_ptr().add(i)),
7538            );
7539            let xb23 = vcombine_s8(
7540                vld1_s8(xs[2].as_ptr().add(i)),
7541                vld1_s8(xs[3].as_ptr().add(i)),
7542            );
7543            asm!(
7544                "smmla {a01:v}.4s, {w:v}.16b, {x01:v}.16b",
7545                "smmla {a23:v}.4s, {w:v}.16b, {x23:v}.16b",
7546                a01 = inout(vreg) acc01, a23 = inout(vreg) acc23,
7547                w = in(vreg) wa, x01 = in(vreg) xb01, x23 = in(vreg) xb23,
7548                options(pure, nomem, nostack),
7549            );
7550            i += 8;
7551        }
7552        let mut out = [[0i32; 4]; 2];
7553        let a01: [i32; 4] = core::mem::transmute(acc01);
7554        let a23: [i32; 4] = core::mem::transmute(acc23);
7555        out[0][0] = a01[0];
7556        out[0][1] = a01[1];
7557        out[1][0] = a01[2];
7558        out[1][1] = a01[3];
7559        out[0][2] = a23[0];
7560        out[0][3] = a23[1];
7561        out[1][2] = a23[2];
7562        out[1][3] = a23[3];
7563        if i < n {
7564            for (k, x) in xs.iter().enumerate() {
7565                for j in i..n {
7566                    out[0][k] += (w0[j] as i8) as i32 * x[j] as i32;
7567                    out[1][k] += (w1[j] as i8) as i32 * x[j] as i32;
7568                }
7569            }
7570        }
7571        out
7572    }
7573}
7574
7575/// ARM twin of the x86 blocked prefill GEMM: two weight rows stay in
7576/// registers across four activation streams, eight sdot accumulators.
7577/// (The per-row form re-read each W row once per activation.)
7578#[cfg(target_arch = "aarch64")]
7579#[target_feature(enable = "neon,dotprod")]
7580unsafe fn dot_i8_sdot_2x4(w0: &[u8], w1: &[u8], xs: [&[i8]; 4]) -> [[i32; 4]; 2] {
7581    // SAFETY: callers uphold slice-length contracts.
7582    unsafe {
7583        use core::arch::aarch64::*;
7584        use core::arch::asm;
7585        let n = w0.len();
7586        let w0p = w0.as_ptr() as *const i8;
7587        let w1p = w1.as_ptr() as *const i8;
7588        let mut acc = [[vdupq_n_s32(0); 4]; 2];
7589        let mut i = 0usize;
7590        while i + 16 <= n {
7591            let wv0 = vld1q_s8(w0p.add(i));
7592            let wv1 = vld1q_s8(w1p.add(i));
7593            for (k, x) in xs.iter().enumerate() {
7594                let xv = vld1q_s8(x.as_ptr().add(i));
7595                let (mut a0, mut a1) = (acc[0][k], acc[1][k]);
7596                asm!(
7597                    "sdot {a0:v}.4s, {w0:v}.16b, {x:v}.16b",
7598                    "sdot {a1:v}.4s, {w1:v}.16b, {x:v}.16b",
7599                    a0 = inout(vreg) a0, a1 = inout(vreg) a1,
7600                    w0 = in(vreg) wv0, w1 = in(vreg) wv1, x = in(vreg) xv,
7601                    options(pure, nomem, nostack),
7602                );
7603                acc[0][k] = a0;
7604                acc[1][k] = a1;
7605            }
7606            i += 16;
7607        }
7608        let mut out = [[0i32; 4]; 2];
7609        for r in 0..2 {
7610            for k in 0..4 {
7611                out[r][k] = vaddvq_s32(acc[r][k]);
7612            }
7613        }
7614        if i < n {
7615            for (k, x) in xs.iter().enumerate() {
7616                for j in i..n {
7617                    out[0][k] += (w0[j] as i8) as i32 * x[j] as i32;
7618                    out[1][k] += (w1[j] as i8) as i32 * x[j] as i32;
7619                }
7620            }
7621        }
7622        out
7623    }
7624}
7625
7626/// Blocked 2 weight rows × 4 activations for the prefill GEMM
7627/// (roadmap P0: packed panels + multi-row accumulators). The two rows'
7628/// abs() live in registers across all four activation streams; the
7629/// sign-fixup is recomputed per pair (the price of the maddubs trick).
7630/// Returns raw i8·i8 dots; the caller applies scales and outliers.
7631#[cfg(target_arch = "x86_64")]
7632#[target_feature(enable = "avx2")]
7633unsafe fn dot_i8_i8_avx2_2x4(w0: &[u8], w1: &[u8], xs: [&[i8]; 4]) -> [[i32; 4]; 2] {
7634    // SAFETY: callers uphold slice-length contracts.
7635    unsafe {
7636        use core::arch::x86_64::*;
7637        let n = w0.len();
7638        let ones = _mm256_set1_epi16(1);
7639        let mut acc = [[_mm256_setzero_si256(); 4]; 2];
7640        let mut j = 0usize;
7641        while j + 32 <= n {
7642            let wv0 = _mm256_loadu_si256(w0.as_ptr().add(j) as *const __m256i);
7643            let wv1 = _mm256_loadu_si256(w1.as_ptr().add(j) as *const __m256i);
7644            let aw0 = _mm256_abs_epi8(wv0);
7645            let aw1 = _mm256_abs_epi8(wv1);
7646            for (k, x) in xs.iter().enumerate() {
7647                let xv = _mm256_loadu_si256(x.as_ptr().add(j) as *const __m256i);
7648                let p0 = _mm256_maddubs_epi16(aw0, _mm256_sign_epi8(xv, wv0));
7649                acc[0][k] = _mm256_add_epi32(acc[0][k], _mm256_madd_epi16(p0, ones));
7650                let p1 = _mm256_maddubs_epi16(aw1, _mm256_sign_epi8(xv, wv1));
7651                acc[1][k] = _mm256_add_epi32(acc[1][k], _mm256_madd_epi16(p1, ones));
7652            }
7653            j += 32;
7654        }
7655        let mut out = [[0i32; 4]; 2];
7656        for r in 0..2 {
7657            for k in 0..4 {
7658                let a = acc[r][k];
7659                let hi128 = _mm256_extracti128_si256::<1>(a);
7660                let s128 = _mm_add_epi32(_mm256_castsi256_si128(a), hi128);
7661                let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
7662                let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
7663                out[r][k] = _mm_cvtsi128_si32(s32);
7664            }
7665        }
7666        if j < n {
7667            for (k, x) in xs.iter().enumerate() {
7668                for i in j..n {
7669                    out[0][k] += (w0[i] as i8) as i32 * x[i] as i32;
7670                    out[1][k] += (w1[i] as i8) as i32 * x[i] as i32;
7671                }
7672            }
7673        }
7674        out
7675    }
7676}
7677
7678/// AVX2/VNNI q8 row dot with exact outlier correction (x86 mirror of
7679/// `row_dot_sdot` — same A8W8 contract). With AVX-512 VNNI the row goes
7680/// through the bias trick: Σ(w+128)·x via pure `vpdpbusd` (no per-lane
7681/// sign fixups), corrected by −128·Σx with Σx precomputed per split.
7682#[cfg(target_arch = "x86_64")]
7683#[inline]
7684fn row_dot_avx2(row: &[u8], act: &SplitAct) -> f32 {
7685    let dot = if avx512vnni_enabled() && row.len() >= 64 {
7686        (unsafe { dot_u8p128_i8_vnni(row, &act.xq) }) - 128 * act.xsum
7687    } else {
7688        unsafe { dot_i8_i8_avx2(row, &act.xq) }
7689    };
7690    let mut acc = dot as f32 * act.sx;
7691    for &(j, xv) in &act.outliers {
7692        acc += (row[j] as i8) as f32 * xv;
7693    }
7694    acc
7695}
7696
7697/// Σ (w[i]+128)·x[i] via pure `vpdpbusd` — the caller subtracts
7698/// 128·Σx. Four independent accumulators (dpbusd is ~5-cycle latency;
7699/// a single-acc loop runs latency-bound, measured on Granite Rapids).
7700#[cfg(target_arch = "x86_64")]
7701#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
7702unsafe fn dot_u8p128_i8_vnni(w: &[u8], xq: &[i8]) -> i32 {
7703    // SAFETY: callers uphold slice-length contracts (see call sites).
7704    unsafe {
7705        use core::arch::x86_64::*;
7706        let n = w.len();
7707        let flip = _mm512_set1_epi8(-128); // XOR 0x80: i8 w → u8 (w+128)
7708        #[inline(always)]
7709        unsafe fn step(
7710            w: *const u8,
7711            x: *const i8,
7712            flip: core::arch::x86_64::__m512i,
7713            acc: core::arch::x86_64::__m512i,
7714        ) -> core::arch::x86_64::__m512i {
7715            unsafe {
7716                use core::arch::x86_64::*;
7717                let wv = _mm512_xor_si512(_mm512_loadu_si512(w as *const _), flip);
7718                _mm512_dpbusd_epi32(acc, wv, _mm512_loadu_si512(x as *const _))
7719            }
7720        }
7721        let (mut a0, mut a1, mut a2, mut a3) = (
7722            _mm512_setzero_si512(),
7723            _mm512_setzero_si512(),
7724            _mm512_setzero_si512(),
7725            _mm512_setzero_si512(),
7726        );
7727        let mut j = 0usize;
7728        while j + 256 <= n {
7729            a0 = step(w.as_ptr().add(j), xq.as_ptr().add(j), flip, a0);
7730            a1 = step(w.as_ptr().add(j + 64), xq.as_ptr().add(j + 64), flip, a1);
7731            a2 = step(w.as_ptr().add(j + 128), xq.as_ptr().add(j + 128), flip, a2);
7732            a3 = step(w.as_ptr().add(j + 192), xq.as_ptr().add(j + 192), flip, a3);
7733            j += 256;
7734        }
7735        while j + 64 <= n {
7736            a0 = step(w.as_ptr().add(j), xq.as_ptr().add(j), flip, a0);
7737            j += 64;
7738        }
7739        let mut total = _mm512_reduce_add_epi32(_mm512_add_epi32(
7740            _mm512_add_epi32(a0, a1),
7741            _mm512_add_epi32(a2, a3),
7742        ));
7743        // Scalar tail: (w as i8) + 128 ≡ (w as u8) ^ 0x80.
7744        while j < n {
7745            total += ((w[j] ^ 0x80) as i32) * xq[j] as i32;
7746            j += 1;
7747        }
7748        total
7749    }
7750}
7751
7752/// One q4 row via AVX2: nibbles → centered i8 (unpacklo/hi restores the
7753/// writer's flat order, same as the NEON vzip pair), maddubs against
7754/// the pre-quantized activation group, × the group's f16 scale. Pair
7755/// saturation safe: |w|≤8, |x|≤127 → 2·8·127 ≪ 32767. Mirror of
7756/// `dot_q4_row_sdot`.
7757#[cfg(target_arch = "x86_64")]
7758#[target_feature(enable = "avx2")]
7759unsafe fn dot_q4_row_avx2(packed: &[u8], scales: &[u8], g0: usize, gpr: usize, xq: &[i8]) -> f32 {
7760    // SAFETY: callers uphold slice-length contracts (16 packed bytes and
7761    // 2 scale bytes per group; xq.len() == gpr·GROUP_SIZE).
7762    unsafe {
7763        use core::arch::x86_64::*;
7764        let lomask = _mm_set1_epi8(0x0F);
7765        let eight = _mm256_set1_epi8(8);
7766        let ones = _mm256_set1_epi16(1);
7767        let mut acc = 0f32;
7768        for gi in 0..gpr {
7769            let g = g0 + gi;
7770            let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
7771            let b = _mm_loadu_si128(packed.as_ptr().add(g * 16) as *const __m128i);
7772            let lo = _mm_and_si128(b, lomask);
7773            let hi = _mm_and_si128(_mm_srli_epi16::<4>(b), lomask);
7774            let w = _mm256_sub_epi8(
7775                _mm256_set_m128i(_mm_unpackhi_epi8(lo, hi), _mm_unpacklo_epi8(lo, hi)),
7776                eight,
7777            );
7778            let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
7779            let p16 = _mm256_maddubs_epi16(_mm256_abs_epi8(w), _mm256_sign_epi8(x, w));
7780            let d = _mm256_madd_epi16(p16, ones);
7781            let hi128 = _mm256_extracti128_si256::<1>(d);
7782            let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
7783            let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
7784            let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
7785            acc += _mm_cvtsi128_si32(s32) as f32 * s;
7786        }
7787        acc
7788    }
7789}
7790
7791/// Two-activation q4 row via AVX2: nibbles unpacked ONCE per group,
7792/// both activations dotted against the same centered i8 register.
7793#[cfg(target_arch = "x86_64")]
7794#[target_feature(enable = "avx2")]
7795unsafe fn dot_q4_row_avx2_2(
7796    packed: &[u8],
7797    scales: &[u8],
7798    g0: usize,
7799    gpr: usize,
7800    xq1: &[i8],
7801    xq2: &[i8],
7802) -> (f32, f32) {
7803    // SAFETY: callers uphold slice-length contracts (see dot_q4_row_avx2).
7804    unsafe {
7805        use core::arch::x86_64::*;
7806        let lomask = _mm_set1_epi8(0x0F);
7807        let eight = _mm256_set1_epi8(8);
7808        let ones = _mm256_set1_epi16(1);
7809        let (mut acc1, mut acc2) = (0f32, 0f32);
7810        #[inline(always)]
7811        unsafe fn hsum(d: core::arch::x86_64::__m256i) -> i32 {
7812            unsafe {
7813                use core::arch::x86_64::*;
7814                let hi128 = _mm256_extracti128_si256::<1>(d);
7815                let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
7816                let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
7817                let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
7818                _mm_cvtsi128_si32(s32)
7819            }
7820        }
7821        for gi in 0..gpr {
7822            let g = g0 + gi;
7823            let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
7824            let b = _mm_loadu_si128(packed.as_ptr().add(g * 16) as *const __m128i);
7825            let lo = _mm_and_si128(b, lomask);
7826            let hi = _mm_and_si128(_mm_srli_epi16::<4>(b), lomask);
7827            let w = _mm256_sub_epi8(
7828                _mm256_set_m128i(_mm_unpackhi_epi8(lo, hi), _mm_unpacklo_epi8(lo, hi)),
7829                eight,
7830            );
7831            let aw = _mm256_abs_epi8(w);
7832            let x1 = _mm256_loadu_si256(xq1.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
7833            let x2 = _mm256_loadu_si256(xq2.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
7834            let d1 = _mm256_madd_epi16(_mm256_maddubs_epi16(aw, _mm256_sign_epi8(x1, w)), ones);
7835            let d2 = _mm256_madd_epi16(_mm256_maddubs_epi16(aw, _mm256_sign_epi8(x2, w)), ones);
7836            acc1 += hsum(d1) as f32 * s;
7837            acc2 += hsum(d2) as f32 * s;
7838        }
7839        (acc1, acc2)
7840    }
7841}
7842
7843/// One q8 row range via AVX2 (x86 mirror of `q8_range_sdot`).
7844#[cfg(target_arch = "x86_64")]
7845fn q8_range_avx2(
7846    q: &[u8],
7847    row_scale: &[f32],
7848    act: &SplitAct,
7849    cols: usize,
7850    out_addr: SendMut,
7851    start: usize,
7852    end: usize,
7853) {
7854    for o in start..end {
7855        let v = row_dot_avx2(&q[o * cols..(o + 1) * cols], act) * row_scale[o];
7856        // SAFETY: disjoint row ranges per worker.
7857        unsafe { *out_addr.at(o) = v };
7858    }
7859}
7860
7861/// Two-input q8 row range via AVX2 (x86 mirror of `q8_range2_sdot`).
7862#[cfg(target_arch = "x86_64")]
7863#[allow(clippy::too_many_arguments)]
7864fn q8_range2_avx2(
7865    q: &[u8],
7866    row_scale: &[f32],
7867    a1: &SplitAct,
7868    a2: &SplitAct,
7869    cols: usize,
7870    p1: SendMut,
7871    p2: SendMut,
7872    start: usize,
7873    end: usize,
7874) {
7875    for o in start..end {
7876        let row = &q[o * cols..(o + 1) * cols];
7877        // SAFETY: disjoint row ranges per worker.
7878        unsafe {
7879            *p1.at(o) = row_dot_avx2(row, a1) * row_scale[o];
7880            *p2.at(o) = row_dot_avx2(row, a2) * row_scale[o];
7881        }
7882    }
7883}
7884
7885// ───────────────────── A8W8 SDOT path (port of vmfcore, ×1.78 decode) ─────────────────────
7886
7887/// ARMv8.6 i8mm (smmla): 32 int8 MACs per instruction vs sdot's 16 —
7888/// yet MEASURED 2.4× SLOWER than the blocked sdot on Apple silicon
7889/// (108 vs 264 GF/s): the on-the-fly vcombine packing and the two-
7890/// accumulator dependency chain swamp the MAC advantage, and Apple's
7891/// four SIMD pipes already keep sdot fed. OPT-IN (CMF_I8MM=1) for
7892/// field trials on Cortex-A710/X-class parts with two pipes, where the
7893/// balance may differ; a pre-interleaved weight layout (repack infra)
7894/// is the known path if it ever earns its keep.
7895#[cfg(target_arch = "aarch64")]
7896fn i8mm_enabled() -> bool {
7897    static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
7898    *ON.get_or_init(|| {
7899        std::env::var("CMF_I8MM").map(|v| v == "1").unwrap_or(false)
7900            && std::arch::is_aarch64_feature_detected!("i8mm")
7901    })
7902}
7903
7904/// SDOT enabled? Default ON when the CPU has ARMv8.2 dotprod;
7905/// `CMF_SDOT=0` disables (falls back to i8×f32 NEON).
7906/// (On non-ARM release builds only the test tolerance switch calls it.)
7907#[cfg_attr(not(target_arch = "aarch64"), allow(dead_code))]
7908fn sdot_enabled() -> bool {
7909    use std::sync::OnceLock;
7910    static ON: OnceLock<bool> = OnceLock::new();
7911    *ON.get_or_init(|| {
7912        let want = std::env::var("CMF_SDOT").map(|v| v != "0").unwrap_or(true);
7913        if !want {
7914            return false;
7915        }
7916
7917        #[cfg(target_arch = "aarch64")]
7918        {
7919            if std::arch::is_aarch64_feature_detected!("dotprod") {
7920                return true;
7921            }
7922            #[cfg(target_os = "android")]
7923            {
7924                if let Ok(cpuinfo) = std::fs::read_to_string("/proc/cpuinfo") {
7925                    if cpuinfo.lines().any(|l| {
7926                        (l.starts_with("Features") || l.starts_with("features"))
7927                            && l.contains("asimddp")
7928                    }) {
7929                        return true;
7930                    }
7931                }
7932            }
7933            false
7934        }
7935        #[cfg(not(target_arch = "aarch64"))]
7936        {
7937            false
7938        }
7939    })
7940}
7941
7942/// Two-field activation split (≡ vmfcore `q8_split_prep`): outlier
7943/// channels (>8·rms) are computed exactly in f32; the bulk (outliers
7944/// zeroed → clean absmax) goes through int8 SDOT. Computed ONCE per
7945/// matvec, shared by all rows/workers.
7946struct SplitAct {
7947    xq: Vec<i8>,
7948    sx: f32,
7949    outliers: Vec<(usize, f32)>,
7950    /// Σ xq — the VNNI bias-trick correction (`(w+128)·x` sums need
7951    /// `−128·Σx`); one i32 per split, computed once per matvec.
7952    #[cfg_attr(not(target_arch = "x86_64"), allow(dead_code))]
7953    xsum: i32,
7954}
7955
7956thread_local! {
7957    /// Recycled xq buffers: split_act runs for every matvec (~200/token)
7958    /// and its hidden-size allocation was steady-state heap churn.
7959    static XQ_FREE: std::cell::RefCell<Vec<Vec<i8>>> =
7960        const { std::cell::RefCell::new(Vec::new()) };
7961}
7962
7963impl Drop for SplitAct {
7964    fn drop(&mut self) {
7965        let buf = std::mem::take(&mut self.xq);
7966        if buf.capacity() > 0 {
7967            XQ_FREE.with(|f| {
7968                let mut f = f.borrow_mut();
7969                if f.len() < 16 {
7970                    f.push(buf);
7971                }
7972            });
7973        }
7974    }
7975}
7976
7977fn split_act(x: &[f32]) -> SplitAct {
7978    let n = x.len();
7979    let rms = (x.iter().map(|&v| (v * v) as f64).sum::<f64>() / n.max(1) as f64).sqrt() as f32;
7980    let thr = 8.0 * rms;
7981    // One pass: collect outliers and the bulk absmax (outliers excluded —
7982    // identical to the old zero-then-fold over a copied buffer, minus the
7983    // full-vector copy).
7984    let mut outliers: Vec<(usize, f32)> = Vec::new();
7985    let mut amax = 0f32;
7986    for (j, &v) in x.iter().enumerate() {
7987        let a = v.abs();
7988        if a > thr {
7989            outliers.push((j, v));
7990        } else if a > amax {
7991            amax = a;
7992        }
7993    }
7994    let sx = if amax > 0.0 { amax / 127.0 } else { 1.0 };
7995    let inv = 1.0 / sx;
7996    let mut xq = XQ_FREE.with(|f| f.borrow_mut().pop()).unwrap_or_default();
7997    xq.clear();
7998    xq.reserve(n);
7999    if outliers.is_empty() {
8000        xq.extend(
8001            x.iter()
8002                .map(|&v| (v * inv).round().clamp(-127.0, 127.0) as i8),
8003        );
8004    } else {
8005        // Outlier slots quantize to 0 (their exact term is added later).
8006        xq.extend(x.iter().map(|&v| {
8007            if v.abs() > thr {
8008                0
8009            } else {
8010                (v * inv).round().clamp(-127.0, 127.0) as i8
8011            }
8012        }));
8013    }
8014    let xsum = xq.iter().map(|&v| v as i32).sum();
8015    SplitAct {
8016        xq,
8017        sx,
8018        outliers,
8019        xsum,
8020    }
8021}
8022
8023fn split_act_q8_2f(x: &[f32], col: &[f32]) -> SplitAct {
8024    let n = x.len();
8025    let rms = (x
8026        .iter()
8027        .zip(col)
8028        .map(|(&a, &c)| {
8029            let v = a * c;
8030            (v * v) as f64
8031        })
8032        .sum::<f64>()
8033        / n.max(1) as f64)
8034        .sqrt() as f32;
8035    let thr = 8.0 * rms;
8036
8037    let mut outliers = Vec::new();
8038    let mut amax = 0f32;
8039    for (j, (&a, &c)) in x.iter().zip(col).enumerate() {
8040        let v = a * c;
8041        let s = v.abs();
8042        if s > thr {
8043            outliers.push((j, v));
8044        } else if s > amax {
8045            amax = s;
8046        }
8047    }
8048
8049    let sx = if amax > 0.0 { amax / 127.0 } else { 1.0 };
8050    let inv = 1.0 / sx;
8051    let mut xq = XQ_FREE.with(|f| f.borrow_mut().pop()).unwrap_or_default();
8052    xq.clear();
8053    xq.reserve(n);
8054    if outliers.is_empty() {
8055        xq.extend(
8056            x.iter()
8057                .zip(col)
8058                .map(|(&a, &c)| ((a * c) * inv).round().clamp(-127.0, 127.0) as i8),
8059        );
8060    } else {
8061        xq.extend(x.iter().zip(col).map(|(&a, &c)| {
8062            let v = a * c;
8063            if v.abs() > thr {
8064                0
8065            } else {
8066                (v * inv).round().clamp(-127.0, 127.0) as i8
8067            }
8068        }));
8069    }
8070    let xsum = xq.iter().map(|&v| v as i32).sum();
8071    SplitAct {
8072        xq,
8073        sx,
8074        outliers,
8075        xsum,
8076    }
8077}
8078
8079/// int8(weight)·int8(activation) → i32 via `sdot` (inline asm — the
8080/// vdotq intrinsic is unstable; port of vmfcore `dot_i8_sdot`).
8081#[cfg(target_arch = "aarch64")]
8082#[target_feature(enable = "neon,dotprod")]
8083unsafe fn dot_i8_sdot(w: &[u8], xq: &[i8]) -> i32 {
8084    // SAFETY: callers uphold slice-length contracts (see call sites).
8085    unsafe {
8086        use core::arch::aarch64::*;
8087        use core::arch::asm;
8088        let wp = w.as_ptr() as *const i8;
8089        let n = w.len();
8090        let (mut a0, mut a1, mut a2, mut a3) = (
8091            vdupq_n_s32(0),
8092            vdupq_n_s32(0),
8093            vdupq_n_s32(0),
8094            vdupq_n_s32(0),
8095        );
8096        let mut i = 0;
8097        while i + 64 <= n {
8098            let (w0, x0) = (vld1q_s8(wp.add(i)), vld1q_s8(xq.as_ptr().add(i)));
8099            let (w1, x1) = (vld1q_s8(wp.add(i + 16)), vld1q_s8(xq.as_ptr().add(i + 16)));
8100            let (w2, x2) = (vld1q_s8(wp.add(i + 32)), vld1q_s8(xq.as_ptr().add(i + 32)));
8101            let (w3, x3) = (vld1q_s8(wp.add(i + 48)), vld1q_s8(xq.as_ptr().add(i + 48)));
8102            asm!(
8103                "sdot {a0:v}.4s, {w0:v}.16b, {x0:v}.16b",
8104                "sdot {a1:v}.4s, {w1:v}.16b, {x1:v}.16b",
8105                "sdot {a2:v}.4s, {w2:v}.16b, {x2:v}.16b",
8106                "sdot {a3:v}.4s, {w3:v}.16b, {x3:v}.16b",
8107                a0 = inout(vreg) a0, a1 = inout(vreg) a1, a2 = inout(vreg) a2, a3 = inout(vreg) a3,
8108                w0 = in(vreg) w0, x0 = in(vreg) x0, w1 = in(vreg) w1, x1 = in(vreg) x1,
8109                w2 = in(vreg) w2, x2 = in(vreg) x2, w3 = in(vreg) w3, x3 = in(vreg) x3,
8110                options(pure, nomem, nostack),
8111            );
8112            i += 64;
8113        }
8114        while i + 16 <= n {
8115            let (wv, xv) = (vld1q_s8(wp.add(i)), vld1q_s8(xq.as_ptr().add(i)));
8116            asm!("sdot {a:v}.4s, {w:v}.16b, {x:v}.16b",
8117                 a = inout(vreg) a0, w = in(vreg) wv, x = in(vreg) xv, options(pure, nomem, nostack));
8118            i += 16;
8119        }
8120        let mut s = vaddvq_s32(vaddq_s32(vaddq_s32(a0, a1), vaddq_s32(a2, a3)));
8121        while i < n {
8122            s += (*wp.add(i)) as i32 * xq[i] as i32;
8123            i += 1;
8124        }
8125        s
8126    }
8127}
8128
8129/// Row-blocked SDOT: 4 output rows per pass — the activation chunk is
8130/// loaded once and reused, 4 independent accumulators hide sdot latency
8131/// (port of vmfcore `dot_i8_sdot_4rows`).
8132#[cfg(target_arch = "aarch64")]
8133#[target_feature(enable = "neon,dotprod")]
8134unsafe fn dot_i8_sdot_4rows(w0: &[u8], w1: &[u8], w2: &[u8], w3: &[u8], xq: &[i8]) -> [i32; 4] {
8135    // SAFETY: callers uphold slice-length contracts (see call sites).
8136    unsafe {
8137        use core::arch::aarch64::*;
8138        use core::arch::asm;
8139        let n = xq.len();
8140        let px = xq.as_ptr();
8141        let (p0, p1, p2, p3) = (
8142            w0.as_ptr() as *const i8,
8143            w1.as_ptr() as *const i8,
8144            w2.as_ptr() as *const i8,
8145            w3.as_ptr() as *const i8,
8146        );
8147        let (mut a0, mut a1, mut a2, mut a3) = (
8148            vdupq_n_s32(0),
8149            vdupq_n_s32(0),
8150            vdupq_n_s32(0),
8151            vdupq_n_s32(0),
8152        );
8153        let mut i = 0;
8154        while i + 16 <= n {
8155            let x = vld1q_s8(px.add(i));
8156            let v0 = vld1q_s8(p0.add(i));
8157            let v1 = vld1q_s8(p1.add(i));
8158            let v2 = vld1q_s8(p2.add(i));
8159            let v3 = vld1q_s8(p3.add(i));
8160            asm!(
8161                "sdot {a0:v}.4s, {v0:v}.16b, {x:v}.16b",
8162                "sdot {a1:v}.4s, {v1:v}.16b, {x:v}.16b",
8163                "sdot {a2:v}.4s, {v2:v}.16b, {x:v}.16b",
8164                "sdot {a3:v}.4s, {v3:v}.16b, {x:v}.16b",
8165                a0 = inout(vreg) a0, a1 = inout(vreg) a1, a2 = inout(vreg) a2, a3 = inout(vreg) a3,
8166                v0 = in(vreg) v0, v1 = in(vreg) v1, v2 = in(vreg) v2, v3 = in(vreg) v3, x = in(vreg) x,
8167                options(pure, nomem, nostack),
8168            );
8169            i += 16;
8170        }
8171        let mut r = [
8172            vaddvq_s32(a0),
8173            vaddvq_s32(a1),
8174            vaddvq_s32(a2),
8175            vaddvq_s32(a3),
8176        ];
8177        while i < n {
8178            let xi = *px.add(i) as i32;
8179            r[0] += (*p0.add(i)) as i32 * xi;
8180            r[1] += (*p1.add(i)) as i32 * xi;
8181            r[2] += (*p2.add(i)) as i32 * xi;
8182            r[3] += (*p3.add(i)) as i32 * xi;
8183            i += 1;
8184        }
8185        r
8186    }
8187}
8188
8189/// 4 interleaved rows in one pass: the repacked group is [r0[c], r1[c],
8190/// r2[c], r3[c]] per 16-byte chunk, so each iteration reads ONE 64-byte
8191/// line plus the shared activation chunk — a single sequential weight
8192/// stream per worker. Per-row accumulation is the same one-accumulator
8193/// scheme as `dot_i8_sdot_4rows`; integer sums are exact, so outputs
8194/// are bit-identical to the mmap-layout kernel.
8195#[cfg(target_arch = "aarch64")]
8196#[target_feature(enable = "neon,dotprod")]
8197unsafe fn dot_i8_sdot_4rows_il(g: &[u8], xq: &[i8]) -> [i32; 4] {
8198    // SAFETY: callers uphold slice-length contracts (g.len() == 4·n,
8199    // n % 16 == 0 — guaranteed by the repack gate).
8200    unsafe {
8201        use core::arch::aarch64::*;
8202        use core::arch::asm;
8203        let n = xq.len();
8204        let px = xq.as_ptr();
8205        let pg = g.as_ptr() as *const i8;
8206        let (mut a0, mut a1, mut a2, mut a3) = (
8207            vdupq_n_s32(0),
8208            vdupq_n_s32(0),
8209            vdupq_n_s32(0),
8210            vdupq_n_s32(0),
8211        );
8212        let mut i = 0;
8213        while i + 16 <= n {
8214            let x = vld1q_s8(px.add(i));
8215            let base = pg.add(4 * i);
8216            let v0 = vld1q_s8(base);
8217            let v1 = vld1q_s8(base.add(16));
8218            let v2 = vld1q_s8(base.add(32));
8219            let v3 = vld1q_s8(base.add(48));
8220            asm!(
8221                "sdot {a0:v}.4s, {v0:v}.16b, {x:v}.16b",
8222                "sdot {a1:v}.4s, {v1:v}.16b, {x:v}.16b",
8223                "sdot {a2:v}.4s, {v2:v}.16b, {x:v}.16b",
8224                "sdot {a3:v}.4s, {v3:v}.16b, {x:v}.16b",
8225                a0 = inout(vreg) a0, a1 = inout(vreg) a1, a2 = inout(vreg) a2, a3 = inout(vreg) a3,
8226                v0 = in(vreg) v0, v1 = in(vreg) v1, v2 = in(vreg) v2, v3 = in(vreg) v3, x = in(vreg) x,
8227                options(pure, nomem, nostack),
8228            );
8229            i += 16;
8230        }
8231        [
8232            vaddvq_s32(a0),
8233            vaddvq_s32(a1),
8234            vaddvq_s32(a2),
8235            vaddvq_s32(a3),
8236        ]
8237    }
8238}
8239
8240/// One q8 row range via SDOT (4-row blocks + tail) — the body of
8241/// `qmatvec`'s hot loop, extracted so multi-matrix jobs can drive the
8242/// SAME kernel for several tensors under one pool dispatch. `rep` — the
8243/// load-time interleaved repack (empty = mmap layout only); rows outside
8244/// full 4-row groups always come from the mmap layout.
8245#[cfg(target_arch = "aarch64")]
8246fn q8_range_sdot(
8247    q: &[u8],
8248    rep: &[u8],
8249    row_scale: &[f32],
8250    act: &SplitAct,
8251    cols: usize,
8252    out_addr: SendMut,
8253    start: usize,
8254    end: usize,
8255) {
8256    let mut o = start;
8257    // Leading rows to the group boundary (repack path only): the pool
8258    // splits row ranges arbitrarily, groups are absolute.
8259    if !rep.is_empty() {
8260        while o < end && o % 4 != 0 {
8261            let v = row_dot_sdot(&q[o * cols..(o + 1) * cols], act) * row_scale[o];
8262            unsafe { *out_addr.at(o) = v };
8263            o += 1;
8264        }
8265    }
8266    while o + 4 <= end {
8267        let r = if rep.is_empty() {
8268            unsafe {
8269                dot_i8_sdot_4rows(
8270                    &q[o * cols..(o + 1) * cols],
8271                    &q[(o + 1) * cols..(o + 2) * cols],
8272                    &q[(o + 2) * cols..(o + 3) * cols],
8273                    &q[(o + 3) * cols..(o + 4) * cols],
8274                    &act.xq,
8275                )
8276            }
8277        } else {
8278            unsafe { dot_i8_sdot_4rows_il(&rep[o * cols..(o + 4) * cols], &act.xq) }
8279        };
8280        for k in 0..4 {
8281            let mut acc = r[k] as f32 * act.sx;
8282            for &(j, xv) in &act.outliers {
8283                acc += (q[(o + k) * cols + j] as i8) as f32 * xv;
8284            }
8285            // SAFETY: disjoint row ranges per worker.
8286            unsafe { *out_addr.at(o + k) = acc * row_scale[o + k] };
8287        }
8288        o += 4;
8289    }
8290    while o < end {
8291        let v = row_dot_sdot(&q[o * cols..(o + 1) * cols], act) * row_scale[o];
8292        unsafe { *out_addr.at(o) = v };
8293        o += 1;
8294    }
8295}
8296
8297/// Two-input q8 row range via SDOT — `qmatvec2`'s hot loop, extracted
8298/// for the fused pair multi-matrix job (`matvec2_many`).
8299#[cfg(target_arch = "aarch64")]
8300#[allow(clippy::too_many_arguments)]
8301fn q8_range2_sdot(
8302    q: &[u8],
8303    row_scale: &[f32],
8304    a1: &SplitAct,
8305    a2: &SplitAct,
8306    cols: usize,
8307    p1: SendMut,
8308    p2: SendMut,
8309    start: usize,
8310    end: usize,
8311) {
8312    for o in start..end {
8313        let row = &q[o * cols..(o + 1) * cols];
8314        // SAFETY: disjoint row ranges per worker.
8315        unsafe {
8316            *p1.at(o) = row_dot_sdot(row, a1) * row_scale[o];
8317            *p2.at(o) = row_dot_sdot(row, a2) * row_scale[o];
8318        }
8319    }
8320}
8321
8322/// Two-input q8 row range, f32 kernel (non-SDOT) — same extraction.
8323#[allow(clippy::too_many_arguments)]
8324fn q8_range2_f32(
8325    q: &[u8],
8326    row_scale: &[f32],
8327    x1: &[f32],
8328    x2: &[f32],
8329    cols: usize,
8330    p1: SendMut,
8331    p2: SendMut,
8332    start: usize,
8333    end: usize,
8334) {
8335    for o in start..end {
8336        let row = &q[o * cols..(o + 1) * cols];
8337        // SAFETY: disjoint row ranges per worker.
8338        unsafe {
8339            *p1.at(o) = dot_i8_f32(row, x1) * row_scale[o];
8340            *p2.at(o) = dot_i8_f32(row, x2) * row_scale[o];
8341        }
8342    }
8343}
8344
8345/// Scalar/NEON-f32 q8 row range (non-SDOT platforms) — same extraction.
8346fn q8_range_f32(
8347    q: &[u8],
8348    row_scale: &[f32],
8349    xs: &[f32],
8350    cols: usize,
8351    out_addr: SendMut,
8352    start: usize,
8353    end: usize,
8354) {
8355    for o in start..end {
8356        let v = dot_i8_f32(&q[o * cols..(o + 1) * cols], xs) * row_scale[o];
8357        // SAFETY: disjoint row ranges per worker.
8358        unsafe { *out_addr.at(o) = v };
8359    }
8360}
8361
8362/// SDOT row dot with exact outlier correction:
8363/// `dot = sdot(w, xq)·sx + Σ_outl w[j]·x[j]` (then × row_scale by caller).
8364#[cfg(target_arch = "aarch64")]
8365#[inline]
8366fn row_dot_sdot(row: &[u8], act: &SplitAct) -> f32 {
8367    let mut acc = unsafe { dot_i8_sdot(row, &act.xq) } as f32 * act.sx;
8368    for &(j, xv) in &act.outliers {
8369        acc += (row[j] as i8) as f32 * xv;
8370    }
8371    acc
8372}
8373
8374/// One q4 row via SDOT: each 32-group's nibbles unpack to centered i8
8375/// (nib−8 ∈ [−8,7]), int8×int8 `sdot` against the pre-quantized
8376/// activation group, × the group's f16 scale. Returns Σ_g dot_g·s_g;
8377/// the caller multiplies by the activation scale and adds the exact
8378/// outlier terms (port of vmfcore `dot_q4_block_sdot`, +23% measured).
8379/// Nibble order matches the writer: element 2k = low nibble, 2k+1 = high
8380/// → zip(lo,hi) restores flat order.
8381#[cfg(target_arch = "aarch64")]
8382#[target_feature(enable = "neon,dotprod")]
8383unsafe fn dot_q4_row_sdot(packed: &[u8], scales: &[u8], g0: usize, gpr: usize, xq: &[i8]) -> f32 {
8384    // SAFETY: callers uphold slice-length contracts (16 packed bytes and
8385    // 2 scale bytes per group; xq.len() == gpr·GROUP_SIZE).
8386    unsafe {
8387        use core::arch::aarch64::*;
8388        use core::arch::asm;
8389        let lomask = vdupq_n_u8(0x0F);
8390        let eight = vdupq_n_s8(8);
8391        let mut acc = 0f32;
8392        for gi in 0..gpr {
8393            let g = g0 + gi;
8394            let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
8395            let b = vld1q_u8(packed.as_ptr().add(g * 16));
8396            let lo = vandq_u8(b, lomask);
8397            let hi = vshrq_n_u8::<4>(b);
8398            let e0 = vsubq_s8(vreinterpretq_s8_u8(vzip1q_u8(lo, hi)), eight);
8399            let e1 = vsubq_s8(vreinterpretq_s8_u8(vzip2q_u8(lo, hi)), eight);
8400            let x0 = vld1q_s8(xq.as_ptr().add(gi * GROUP_SIZE));
8401            let x1 = vld1q_s8(xq.as_ptr().add(gi * GROUP_SIZE + 16));
8402            let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
8403            asm!(
8404                "sdot {a0:v}.4s, {e0:v}.16b, {x0:v}.16b",
8405                "sdot {a1:v}.4s, {e1:v}.16b, {x1:v}.16b",
8406                a0 = inout(vreg) a0, a1 = inout(vreg) a1,
8407                e0 = in(vreg) e0, x0 = in(vreg) x0, e1 = in(vreg) e1, x1 = in(vreg) x1,
8408                options(pure, nomem, nostack),
8409            );
8410            acc += vaddvq_s32(vaddq_s32(a0, a1)) as f32 * s;
8411        }
8412        acc
8413    }
8414}
8415
8416/// Two-activation q4 row via SDOT: the nibble unpack (the expensive
8417/// part) happens ONCE per group; both pre-quantized activations are
8418/// dotted against the same centered i8 registers. Per-lane math matches
8419/// `dot_q4_row_sdot` exactly.
8420#[cfg(target_arch = "aarch64")]
8421#[target_feature(enable = "neon,dotprod")]
8422unsafe fn dot_q4_row_sdot2(
8423    packed: &[u8],
8424    scales: &[u8],
8425    g0: usize,
8426    gpr: usize,
8427    xq1: &[i8],
8428    xq2: &[i8],
8429) -> (f32, f32) {
8430    // SAFETY: callers uphold slice-length contracts (16 packed bytes and
8431    // 2 scale bytes per group; xq*.len() == gpr·GROUP_SIZE).
8432    unsafe {
8433        use core::arch::aarch64::*;
8434        use core::arch::asm;
8435        let lomask = vdupq_n_u8(0x0F);
8436        let eight = vdupq_n_s8(8);
8437        let (mut acc1, mut acc2) = (0f32, 0f32);
8438        for gi in 0..gpr {
8439            let g = g0 + gi;
8440            let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
8441            let b = vld1q_u8(packed.as_ptr().add(g * 16));
8442            let lo = vandq_u8(b, lomask);
8443            let hi = vshrq_n_u8::<4>(b);
8444            let e0 = vsubq_s8(vreinterpretq_s8_u8(vzip1q_u8(lo, hi)), eight);
8445            let e1 = vsubq_s8(vreinterpretq_s8_u8(vzip2q_u8(lo, hi)), eight);
8446            let x10 = vld1q_s8(xq1.as_ptr().add(gi * GROUP_SIZE));
8447            let x11 = vld1q_s8(xq1.as_ptr().add(gi * GROUP_SIZE + 16));
8448            let x20 = vld1q_s8(xq2.as_ptr().add(gi * GROUP_SIZE));
8449            let x21 = vld1q_s8(xq2.as_ptr().add(gi * GROUP_SIZE + 16));
8450            let (mut a0, mut a1, mut b0, mut b1) = (
8451                vdupq_n_s32(0),
8452                vdupq_n_s32(0),
8453                vdupq_n_s32(0),
8454                vdupq_n_s32(0),
8455            );
8456            asm!(
8457                "sdot {a0:v}.4s, {e0:v}.16b, {x10:v}.16b",
8458                "sdot {a1:v}.4s, {e1:v}.16b, {x11:v}.16b",
8459                "sdot {b0:v}.4s, {e0:v}.16b, {x20:v}.16b",
8460                "sdot {b1:v}.4s, {e1:v}.16b, {x21:v}.16b",
8461                a0 = inout(vreg) a0, a1 = inout(vreg) a1,
8462                b0 = inout(vreg) b0, b1 = inout(vreg) b1,
8463                e0 = in(vreg) e0, e1 = in(vreg) e1,
8464                x10 = in(vreg) x10, x11 = in(vreg) x11,
8465                x20 = in(vreg) x20, x21 = in(vreg) x21,
8466                options(pure, nomem, nostack),
8467            );
8468            acc1 += vaddvq_s32(vaddq_s32(a0, a1)) as f32 * s;
8469            acc2 += vaddvq_s32(vaddq_s32(b0, b1)) as f32 * s;
8470        }
8471        (acc1, acc2)
8472    }
8473}
8474
8475// ───────────────────── fused int8 kernels ─────────────────────
8476
8477/// `acc += w · row` where the row is centered i8 — NEON widen+fma on
8478/// aarch64, scalar elsewhere. The KV-cache q8 value path rides on this.
8479#[inline]
8480pub(crate) fn axpy_i8_f32(acc: &mut [f32], row: &[i8], w: f32) {
8481    #[cfg(target_arch = "aarch64")]
8482    unsafe {
8483        return axpy_i8_f32_neon(acc, row, w);
8484    }
8485    #[cfg(target_arch = "x86_64")]
8486    if avx2_enabled() {
8487        return unsafe { axpy_i8_f32_avx2(acc, row, w) };
8488    }
8489    #[allow(unreachable_code)]
8490    {
8491        for (a, &b) in acc.iter_mut().zip(row) {
8492            *a += w * b as f32;
8493        }
8494    }
8495}
8496
8497/// i8→f32 axpy via AVX2/FMA (x86 mirror of `axpy_i8_f32_neon`).
8498#[cfg(target_arch = "x86_64")]
8499#[target_feature(enable = "avx2,fma")]
8500unsafe fn axpy_i8_f32_avx2(acc: &mut [f32], row: &[i8], w: f32) {
8501    // SAFETY: callers uphold slice-length contracts (see call sites).
8502    unsafe {
8503        use core::arch::x86_64::*;
8504        let n = acc.len().min(row.len());
8505        let ap = acc.as_mut_ptr();
8506        let rp = row.as_ptr();
8507        let wv = _mm256_set1_ps(w);
8508        let mut j = 0usize;
8509        while j + 16 <= n {
8510            let rb = _mm_loadu_si128(rp.add(j) as *const __m128i);
8511            let lo = _mm256_cvtepi32_ps(_mm256_cvtepi8_epi32(rb));
8512            let hi = _mm256_cvtepi32_ps(_mm256_cvtepi8_epi32(_mm_srli_si128::<8>(rb)));
8513            let v0 = _mm256_fmadd_ps(wv, lo, _mm256_loadu_ps(ap.add(j)));
8514            let v1 = _mm256_fmadd_ps(wv, hi, _mm256_loadu_ps(ap.add(j + 8)));
8515            _mm256_storeu_ps(ap.add(j), v0);
8516            _mm256_storeu_ps(ap.add(j + 8), v1);
8517            j += 16;
8518        }
8519        while j < n {
8520            *ap.add(j) += w * (*rp.add(j)) as f32;
8521            j += 1;
8522        }
8523    }
8524}
8525
8526#[cfg(target_arch = "aarch64")]
8527#[target_feature(enable = "neon")]
8528unsafe fn axpy_i8_f32_neon(acc: &mut [f32], row: &[i8], w: f32) {
8529    // SAFETY: callers uphold slice-length contracts (see call sites).
8530    unsafe {
8531        use core::arch::aarch64::*;
8532        let n = acc.len().min(row.len());
8533        let ap = acc.as_mut_ptr();
8534        let rp = row.as_ptr();
8535        let wv = vdupq_n_f32(w);
8536        let mut j = 0usize;
8537        while j + 16 <= n {
8538            let rb = vld1q_s8(rp.add(j));
8539            let lo = vmovl_s8(vget_low_s8(rb));
8540            let hi = vmovl_s8(vget_high_s8(rb));
8541            for (off, half) in [(0, lo), (8, hi)] {
8542                let f0 = vcvtq_f32_s32(vmovl_s16(vget_low_s16(half)));
8543                let f1 = vcvtq_f32_s32(vmovl_s16(vget_high_s16(half)));
8544                let o = j + off;
8545                vst1q_f32(ap.add(o), vfmaq_f32(vld1q_f32(ap.add(o)), wv, f0));
8546                vst1q_f32(ap.add(o + 4), vfmaq_f32(vld1q_f32(ap.add(o + 4)), wv, f1));
8547            }
8548            j += 16;
8549        }
8550        while j < n {
8551            *ap.add(j) += w * (*rp.add(j)) as f32;
8552            j += 1;
8553        }
8554    }
8555}
8556
8557/// i8 row · f32 x. NEON on aarch64 (ported from vmfcore `dot_i8_f32_neon`,
8558/// ≈9× scalar), scalar elsewhere.
8559#[inline]
8560pub(crate) fn dot_i8_f32(w: &[u8], x: &[f32]) -> f32 {
8561    #[cfg(target_arch = "aarch64")]
8562    unsafe {
8563        return dot_i8_f32_neon(w, x);
8564    }
8565    #[cfg(target_arch = "x86_64")]
8566    if avx2_enabled() {
8567        return unsafe { dot_i8_f32_avx2(w, x) };
8568    }
8569    #[allow(unreachable_code)]
8570    {
8571        let mut sum = 0.0f32;
8572        for (j, &b) in w.iter().enumerate() {
8573            sum += (b as i8) as f32 * x[j];
8574        }
8575        sum
8576    }
8577}
8578
8579/// i8 row · (x ⊙ col_field) — the q8_2f row dot with the θ col-field
8580/// folded into the product (no prescaled copy of x). NEON on aarch64,
8581/// scalar elsewhere. Used by the active-neuron path `row_dot`.
8582#[inline]
8583fn dot_i8_col_f32(w: &[u8], x: &[f32], col: &[f32]) -> f32 {
8584    #[cfg(target_arch = "aarch64")]
8585    unsafe {
8586        return dot_i8_col_f32_neon(w, x, col);
8587    }
8588    #[allow(unreachable_code)]
8589    {
8590        let mut sum = 0.0f32;
8591        for (j, &b) in w.iter().enumerate() {
8592            sum += (b as i8) as f32 * x[j] * col[j];
8593        }
8594        sum
8595    }
8596}
8597
8598#[cfg(target_arch = "aarch64")]
8599#[target_feature(enable = "neon")]
8600unsafe fn dot_i8_col_f32_neon(w: &[u8], x: &[f32], col: &[f32]) -> f32 {
8601    // SAFETY: callers uphold slice-length contracts (see call sites).
8602    unsafe {
8603        use core::arch::aarch64::*;
8604        let n = x.len();
8605        let wp = w.as_ptr() as *const i8;
8606        let xp = x.as_ptr();
8607        let cp = col.as_ptr();
8608        let (mut a0, mut a1, mut a2, mut a3) = (
8609            vdupq_n_f32(0.0),
8610            vdupq_n_f32(0.0),
8611            vdupq_n_f32(0.0),
8612            vdupq_n_f32(0.0),
8613        );
8614        let mut j = 0usize;
8615        while j + 16 <= n {
8616            let wb = vld1q_s8(wp.add(j));
8617            let lo = vmovl_s8(vget_low_s8(wb));
8618            let hi = vmovl_s8(vget_high_s8(wb));
8619            let w0 = vcvtq_f32_s32(vmovl_s16(vget_low_s16(lo)));
8620            let w1 = vcvtq_f32_s32(vmovl_s16(vget_high_s16(lo)));
8621            let w2 = vcvtq_f32_s32(vmovl_s16(vget_low_s16(hi)));
8622            let w3 = vcvtq_f32_s32(vmovl_s16(vget_high_s16(hi)));
8623            a0 = vfmaq_f32(
8624                a0,
8625                w0,
8626                vmulq_f32(vld1q_f32(xp.add(j)), vld1q_f32(cp.add(j))),
8627            );
8628            a1 = vfmaq_f32(
8629                a1,
8630                w1,
8631                vmulq_f32(vld1q_f32(xp.add(j + 4)), vld1q_f32(cp.add(j + 4))),
8632            );
8633            a2 = vfmaq_f32(
8634                a2,
8635                w2,
8636                vmulq_f32(vld1q_f32(xp.add(j + 8)), vld1q_f32(cp.add(j + 8))),
8637            );
8638            a3 = vfmaq_f32(
8639                a3,
8640                w3,
8641                vmulq_f32(vld1q_f32(xp.add(j + 12)), vld1q_f32(cp.add(j + 12))),
8642            );
8643            j += 16;
8644        }
8645        let mut sum = vaddvq_f32(vaddq_f32(vaddq_f32(a0, a1), vaddq_f32(a2, a3)));
8646        while j < n {
8647            sum += (*wp.add(j)) as f32 * *xp.add(j) * *cp.add(j);
8648            j += 1;
8649        }
8650        sum
8651    }
8652}
8653
8654#[cfg(target_arch = "aarch64")]
8655#[target_feature(enable = "neon")]
8656unsafe fn dot_i8_f32_neon(w: &[u8], x: &[f32]) -> f32 {
8657    // SAFETY: callers uphold slice-length contracts (see call sites).
8658    unsafe {
8659        use core::arch::aarch64::*;
8660        let n = x.len();
8661        let wp = w.as_ptr() as *const i8;
8662        let xp = x.as_ptr();
8663        let (mut a0, mut a1, mut a2, mut a3) = (
8664            vdupq_n_f32(0.0),
8665            vdupq_n_f32(0.0),
8666            vdupq_n_f32(0.0),
8667            vdupq_n_f32(0.0),
8668        );
8669        let mut j = 0usize;
8670        while j + 16 <= n {
8671            let wb = vld1q_s8(wp.add(j));
8672            let lo = vmovl_s8(vget_low_s8(wb));
8673            let hi = vmovl_s8(vget_high_s8(wb));
8674            let w0 = vcvtq_f32_s32(vmovl_s16(vget_low_s16(lo)));
8675            let w1 = vcvtq_f32_s32(vmovl_s16(vget_high_s16(lo)));
8676            let w2 = vcvtq_f32_s32(vmovl_s16(vget_low_s16(hi)));
8677            let w3 = vcvtq_f32_s32(vmovl_s16(vget_high_s16(hi)));
8678            a0 = vfmaq_f32(a0, w0, vld1q_f32(xp.add(j)));
8679            a1 = vfmaq_f32(a1, w1, vld1q_f32(xp.add(j + 4)));
8680            a2 = vfmaq_f32(a2, w2, vld1q_f32(xp.add(j + 8)));
8681            a3 = vfmaq_f32(a3, w3, vld1q_f32(xp.add(j + 12)));
8682            j += 16;
8683        }
8684        let mut sum = vaddvq_f32(vaddq_f32(vaddq_f32(a0, a1), vaddq_f32(a2, a3)));
8685        while j < n {
8686            sum += (*wp.add(j)) as f32 * *xp.add(j);
8687            j += 1;
8688        }
8689        sum
8690    }
8691}
8692
8693#[allow(clippy::too_many_arguments)]
8694fn qmatvec(
8695    q: &[u8],
8696    rep: &[u8],
8697    row_scale: &[f32],
8698    x: &[f32],
8699    col_field: &[f32],
8700    dtype: TensorDtype,
8701    rows: usize,
8702    cols: usize,
8703    out: &mut [f32],
8704    pool: Option<&Pool>,
8705) {
8706    debug_assert_eq!(out.len(), rows);
8707    #[cfg(not(target_arch = "aarch64"))]
8708    let _ = rep;
8709
8710    #[cfg(target_arch = "aarch64")]
8711    if sdot_enabled() {
8712        let act = if dtype == TensorDtype::Q8_2f {
8713            split_act_q8_2f(x, col_field)
8714        } else {
8715            split_act(x)
8716        };
8717        let out_addr = SendMut(out.as_mut_ptr());
8718        let run_range = |start: usize, end: usize| {
8719            q8_range_sdot(q, rep, row_scale, &act, cols, out_addr, start, end)
8720        };
8721        match pool {
8722            Some(pool) if rows >= 256 => pool.run_rows(rows, &run_range),
8723            _ => run_range(0, rows),
8724        }
8725        return;
8726    }
8727    // x86 A8W8 via AVX2 maddubs — same quantized-activation contract as
8728    // the SDOT path (CMF_AVX2=0 keeps the exact i8×f32 loop).
8729    #[cfg(target_arch = "x86_64")]
8730    if avx2_a8w8_enabled() {
8731        let act = if dtype == TensorDtype::Q8_2f {
8732            split_act_q8_2f(x, col_field)
8733        } else {
8734            split_act(x)
8735        };
8736        let out_addr = SendMut(out.as_mut_ptr());
8737        let run_range = |start: usize, end: usize| {
8738            q8_range_avx2(q, row_scale, &act, cols, out_addr, start, end)
8739        };
8740        match pool {
8741            Some(pool) if rows >= 256 => pool.run_rows(rows, &run_range),
8742            _ => run_range(0, rows),
8743        }
8744        return;
8745    }
8746
8747    prescale_with(x, col_field, dtype, 1, |xs| {
8748        let out_addr = SendMut(out.as_mut_ptr());
8749        let run_range = move |start: usize, end: usize| {
8750            for o in start..end {
8751                let v = dot_i8_f32(&q[o * cols..(o + 1) * cols], xs) * row_scale[o];
8752                // SAFETY: disjoint row ranges per worker.
8753                unsafe { *out_addr.at(o) = v };
8754            }
8755        };
8756        match pool {
8757            Some(pool) if rows >= 256 => pool.run_rows(rows, &run_range),
8758            _ => run_range(0, rows),
8759        }
8760    });
8761}
8762
8763#[allow(clippy::too_many_arguments)]
8764fn qmatvec2(
8765    q: &[u8],
8766    row_scale: &[f32],
8767    x1: &[f32],
8768    x2: &[f32],
8769    col_field: &[f32],
8770    dtype: TensorDtype,
8771    rows: usize,
8772    cols: usize,
8773    o1: &mut [f32],
8774    o2: &mut [f32],
8775    pool: Option<&Pool>,
8776) {
8777    #[cfg(target_arch = "aarch64")]
8778    if sdot_enabled() {
8779        let a1s = if dtype == TensorDtype::Q8_2f {
8780            split_act_q8_2f(x1, col_field)
8781        } else {
8782            split_act(x1)
8783        };
8784        let a2s = if dtype == TensorDtype::Q8_2f {
8785            split_act_q8_2f(x2, col_field)
8786        } else {
8787            split_act(x2)
8788        };
8789        let p1 = SendMut(o1.as_mut_ptr());
8790        let p2 = SendMut(o2.as_mut_ptr());
8791        let run_range = |start: usize, end: usize| {
8792            q8_range2_sdot(q, row_scale, &a1s, &a2s, cols, p1, p2, start, end)
8793        };
8794        match pool {
8795            Some(pool) if rows >= 256 => pool.run_rows(rows, &run_range),
8796            _ => run_range(0, rows),
8797        }
8798        return;
8799    }
8800    #[cfg(target_arch = "x86_64")]
8801    if avx2_a8w8_enabled() {
8802        let a1s = if dtype == TensorDtype::Q8_2f {
8803            split_act_q8_2f(x1, col_field)
8804        } else {
8805            split_act(x1)
8806        };
8807        let a2s = if dtype == TensorDtype::Q8_2f {
8808            split_act_q8_2f(x2, col_field)
8809        } else {
8810            split_act(x2)
8811        };
8812        let p1 = SendMut(o1.as_mut_ptr());
8813        let p2 = SendMut(o2.as_mut_ptr());
8814        let run_range = |start: usize, end: usize| {
8815            q8_range2_avx2(q, row_scale, &a1s, &a2s, cols, p1, p2, start, end)
8816        };
8817        match pool {
8818            Some(pool) if rows >= 256 => pool.run_rows(rows, &run_range),
8819            _ => run_range(0, rows),
8820        }
8821        return;
8822    }
8823
8824    prescale_with(x1, col_field, dtype, 1, |x1s| {
8825        prescale_with(x2, col_field, dtype, 2, |x2s| {
8826            let p1 = SendMut(o1.as_mut_ptr());
8827            let p2 = SendMut(o2.as_mut_ptr());
8828            let run_range = move |start: usize, end: usize| {
8829                for o in start..end {
8830                    let row = &q[o * cols..(o + 1) * cols];
8831                    let s1 = dot_i8_f32(row, x1s) * row_scale[o];
8832                    let s2 = dot_i8_f32(row, x2s) * row_scale[o];
8833                    // SAFETY: disjoint row ranges per worker.
8834                    unsafe {
8835                        *p1.at(o) = s1;
8836                        *p2.at(o) = s2;
8837                    }
8838                }
8839            };
8840            match pool {
8841                Some(pool) if rows >= 256 => pool.run_rows(rows, &run_range),
8842                _ => run_range(0, rows),
8843            }
8844        });
8845    });
8846}
8847
8848#[derive(Clone, Copy)]
8849struct SendMut(*mut f32);
8850unsafe impl Send for SendMut {}
8851unsafe impl Sync for SendMut {}
8852
8853impl SendMut {
8854    #[inline]
8855    fn at(self, i: usize) -> *mut f32 {
8856        unsafe { self.0.add(i) }
8857    }
8858}
8859
8860#[cfg(test)]
8861mod tests {
8862    use super::*;
8863
8864    #[test]
8865    fn f32_matvec_matches_matvec_rows_bitexact() {
8866        let (rows, cols) = (300, 40);
8867        let w: Vec<f32> = (0..rows * cols).map(|i| (i as f32 * 0.017).sin()).collect();
8868        let x: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.05).cos()).collect();
8869        let qt = QTensor::from_f32(w.clone(), rows, cols);
8870
8871        let mut a = vec![0.0f32; rows];
8872        matvec_rows(None, &w, &x, &mut a);
8873        let mut b = vec![0.0f32; rows];
8874        qt.matvec(&x, &mut b, None);
8875        assert_eq!(a, b);
8876    }
8877
8878    #[test]
8879    fn sdot_kernel_exact_on_grid() {
8880        // Activations already on the i8 grid (±1 with amax=1 → sx=1/127,
8881        // xq=±127 dequantizes EXACTLY) → the SDOT path must match the
8882        // exact f32 dot to float rounding. This isolates kernel
8883        // correctness from quantization noise.
8884        eprintln!("sdot_enabled = {}", sdot_enabled());
8885        let (rows, cols) = (9, 80); // odd rows → exercises 4-row + tail
8886        let w: Vec<u8> = (0..rows * cols)
8887            .map(|i| (((i * 37) % 251) as i32 - 125) as i8 as u8)
8888            .collect();
8889        let scales: Vec<f32> = (0..rows).map(|o| 0.005 + o as f32 * 0.001).collect();
8890        let x: Vec<f32> = (0..cols)
8891            .map(|i| match i % 3 {
8892                0 => 1.0,
8893                1 => -1.0,
8894                _ => 0.0,
8895            })
8896            .collect();
8897        let mut a = vec![0.0f32; rows];
8898        qmatvec(
8899            &w,
8900            &[],
8901            &scales,
8902            &x,
8903            &[],
8904            TensorDtype::Q8Row,
8905            rows,
8906            cols,
8907            &mut a,
8908            None,
8909        );
8910        for o in 0..rows {
8911            let mut acc = 0.0f32;
8912            for j in 0..cols {
8913                acc += (w[o * cols + j] as i8) as f32 * x[j];
8914            }
8915            let expect = acc * scales[o];
8916            assert!(
8917                (a[o] - expect).abs() < 1e-3 * expect.abs().max(1e-3),
8918                "row {o}: {} vs {expect}",
8919                a[o]
8920            );
8921        }
8922    }
8923
8924    #[test]
8925    fn q1_tbl_fast_path_matches_reference() {
8926        // gpr = 8 exercises the TBL pair-load fast loop, and the LAST
8927        // row's final 4-tile window trips the 4B-overread guard (the
8928        // payload ends exactly at the last tile) — both paths must
8929        // agree with the dequant reference.
8930        let (rows, cols) = (5, 256);
8931        let gpr = cols / GROUP_SIZE;
8932        let mut bytes = Vec::new();
8933        for t in 0..rows * gpr {
8934            let s = 0.007 + (t % 11) as f32 * 0.004;
8935            bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
8936            for j in 0..4 {
8937                bytes.push(((t * 53 + j * 89 + 7) % 249) as u8);
8938            }
8939        }
8940        let x: Vec<f32> = (0..cols)
8941            .map(|i| if (i * 5) % 7 < 3 { 1.0 } else { -1.0 })
8942            .collect();
8943        let mut w = vec![0.0f32; rows * cols];
8944        cortiq_core::quant::dequant_q1(&bytes, &mut w);
8945        let mut got = vec![0.0f32; rows];
8946        q1_matvec(&bytes, &x, rows, cols, &mut got, None);
8947        for o in 0..rows {
8948            let expect: f32 = (0..cols).map(|j| w[o * cols + j] * x[j]).sum();
8949            assert!(
8950                (got[o] - expect).abs() < 1e-3 * expect.abs().max(1e-3),
8951                "row {o}: {} vs {expect}",
8952                got[o]
8953            );
8954        }
8955        // Blocked 1×4 batch (b=5: one quad + remainder) must equal the
8956        // single-matvec path bit-for-bit.
8957        let b = 5usize;
8958        let mut xs_all = Vec::new();
8959        for bi in 0..b {
8960            xs_all.extend(x.iter().map(|v| if bi % 2 == 0 { *v } else { -*v }));
8961        }
8962        let mut mm = vec![0.0f32; b * rows];
8963        q1_matmat(&bytes, &xs_all, b, rows, cols, &mut mm, None);
8964        for bi in 0..b {
8965            let mut single = vec![0.0f32; rows];
8966            q1_matvec(
8967                &bytes,
8968                &xs_all[bi * cols..(bi + 1) * cols],
8969                rows,
8970                cols,
8971                &mut single,
8972                None,
8973            );
8974            assert_eq!(&mm[bi * rows..(bi + 1) * rows], &single[..], "stream {bi}");
8975        }
8976    }
8977
8978    #[test]
8979    fn q1_kernels_match_exact_reference() {
8980        // Synthetic q1 payload: 6-byte tiles [f16 scale][4B bits].
8981        let (rows, cols) = (7, 96);
8982        let gpr = cols / GROUP_SIZE;
8983        let mut bytes = Vec::new();
8984        for t in 0..rows * gpr {
8985            let s = 0.01 + (t % 13) as f32 * 0.003;
8986            bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
8987            for j in 0..4 {
8988                bytes.push(((t * 31 + j * 97) % 251) as u8);
8989            }
8990        }
8991        // On-grid activations (±1, amax 1) → the SDOT path is exact.
8992        let x: Vec<f32> = (0..cols)
8993            .map(|i| if i % 3 == 0 { 1.0 } else { -1.0 })
8994            .collect();
8995        // Reference through the core dequant.
8996        let mut w = vec![0.0f32; rows * cols];
8997        cortiq_core::quant::dequant_q1(&bytes, &mut w);
8998        let mut expect = vec![0.0f32; rows];
8999        for o in 0..rows {
9000            expect[o] = (0..cols).map(|j| w[o * cols + j] * x[j]).sum();
9001        }
9002        let mut got = vec![0.0f32; rows];
9003        q1_matvec(&bytes, &x, rows, cols, &mut got, None);
9004        for o in 0..rows {
9005            assert!(
9006                (got[o] - expect[o]).abs() < 1e-3 * expect[o].abs().max(1e-3),
9007                "row {o}: {} vs {}",
9008                got[o],
9009                expect[o]
9010            );
9011        }
9012        // Pair and batch paths agree with the single path.
9013        let x2: Vec<f32> = x.iter().map(|v| -v).collect();
9014        let (mut a1, mut a2) = (vec![0.0f32; rows], vec![0.0f32; rows]);
9015        q1_matvec2(&bytes, &x, &x2, rows, cols, &mut a1, &mut a2, None);
9016        assert_eq!(a1, got);
9017        let mut xs = x.clone();
9018        xs.extend_from_slice(&x2);
9019        let mut mm = vec![0.0f32; 2 * rows];
9020        q1_matmat(&bytes, &xs, 2, rows, cols, &mut mm, None);
9021        assert_eq!(&mm[..rows], got.as_slice());
9022        assert_eq!(&mm[rows..], a2.as_slice());
9023    }
9024
9025    #[test]
9026    fn repack_is_bit_identical() {
9027        // The interleaved-repack kernel must produce EXACTLY the same
9028        // bits as the mmap-layout kernel: integer accumulation is order-
9029        // exact, the f32 epilogue is identical. Odd rows exercise the
9030        // tail; direct range calls exercise unaligned pool splits.
9031        let (rows, cols) = (267, 96); // 66 groups + 3 tail rows, cols % 16 == 0
9032        let w: Vec<u8> = (0..rows * cols)
9033            .map(|i| (((i * 89) % 253) as i32 - 126) as i8 as u8)
9034            .collect();
9035        let scales: Vec<f32> = (0..rows).map(|o| 0.003 + o as f32 * 0.0007).collect();
9036        let x: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.37).sin() * 2.0).collect();
9037        let rep = q8_repack_layout(&w, rows, cols);
9038        // Group interleave round-trips.
9039        for g in 0..rows / 4 {
9040            for c in 0..cols / 16 {
9041                for lane in 0..4 {
9042                    assert_eq!(
9043                        &rep[g * 4 * cols + c * 64 + lane * 16
9044                            ..g * 4 * cols + c * 64 + lane * 16 + 16],
9045                        &w[(g * 4 + lane) * cols + c * 16..(g * 4 + lane) * cols + c * 16 + 16],
9046                    );
9047                }
9048            }
9049        }
9050        let mut a = vec![0.0f32; rows];
9051        qmatvec(
9052            &w,
9053            &[],
9054            &scales,
9055            &x,
9056            &[],
9057            TensorDtype::Q8Row,
9058            rows,
9059            cols,
9060            &mut a,
9061            None,
9062        );
9063        let mut b = vec![0.0f32; rows];
9064        qmatvec(
9065            &w,
9066            &rep,
9067            &scales,
9068            &x,
9069            &[],
9070            TensorDtype::Q8Row,
9071            rows,
9072            cols,
9073            &mut b,
9074            None,
9075        );
9076        assert_eq!(a, b, "full-range repack output diverged");
9077
9078        #[cfg(target_arch = "aarch64")]
9079        if sdot_enabled() {
9080            // Unaligned range split (pool workers get arbitrary bounds).
9081            let act = split_act(&x);
9082            let mut c1 = vec![0.0f32; rows];
9083            let mut c2 = vec![0.0f32; rows];
9084            q8_range_sdot(
9085                &w,
9086                &[],
9087                &scales,
9088                &act,
9089                cols,
9090                SendMut(c1.as_mut_ptr()),
9091                3,
9092                rows - 2,
9093            );
9094            q8_range_sdot(
9095                &w,
9096                &rep,
9097                &scales,
9098                &act,
9099                cols,
9100                SendMut(c2.as_mut_ptr()),
9101                3,
9102                rows - 2,
9103            );
9104            assert_eq!(c1, c2, "unaligned-range repack output diverged");
9105        }
9106    }
9107
9108    #[test]
9109    fn sdot_a8w8_noise_is_bounded() {
9110        // Off-grid activations: A8 quantization noise must stay small in
9111        // relative L2 over the whole output (realistic accuracy contract;
9112        // vmfcore measured argmax-identical decode on real models).
9113        let (rows, cols) = (16, 512);
9114        let w: Vec<u8> = (0..rows * cols)
9115            .map(|i| (((i * 37) % 251) as i32 - 125) as i8 as u8)
9116            .collect();
9117        let scales = vec![0.01f32; rows];
9118        let x: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.21).sin()).collect();
9119        let mut a = vec![0.0f32; rows];
9120        qmatvec(
9121            &w,
9122            &[],
9123            &scales,
9124            &x,
9125            &[],
9126            TensorDtype::Q8Row,
9127            rows,
9128            cols,
9129            &mut a,
9130            None,
9131        );
9132        let (mut num, mut den) = (0f64, 0f64);
9133        for o in 0..rows {
9134            let mut acc = 0.0f32;
9135            for j in 0..cols {
9136                acc += (w[o * cols + j] as i8) as f32 * x[j];
9137            }
9138            let expect = acc * scales[o];
9139            num += ((a[o] - expect) as f64).powi(2);
9140            den += (expect as f64).powi(2);
9141        }
9142        let rel = (num / den.max(1e-12)).sqrt();
9143        assert!(rel < 0.05, "A8W8 relative L2 error too high: {rel}");
9144    }
9145
9146    #[test]
9147    fn i8_dot_neon_matches_scalar() {
9148        let n = 100;
9149        let w: Vec<u8> = (0..n).map(|i| ((i * 37 + 11) % 251) as u8).collect();
9150        let x: Vec<f32> = (0..n).map(|i| (i as f32 * 0.13).sin()).collect();
9151        let mut scalar = 0.0f32;
9152        for j in 0..n {
9153            scalar += (w[j] as i8) as f32 * x[j];
9154        }
9155        let fast = dot_i8_f32(&w, &x);
9156        assert!((scalar - fast).abs() < 1e-3 * scalar.abs().max(1.0));
9157    }
9158
9159    /// Fused vbit matvec must match full dequant_vbit + dense matvec.
9160    #[test]
9161    fn vbitmatvec_matches_full_dequant() {
9162        let (rows, cols) = (6, 64);
9163        let ng = cols / GROUP_SIZE;
9164        // Hand-craft: bits per row, f16 scales, packed rows.
9165        let bits: Vec<u8> = vec![3, 4, 5, 6, 8, 4];
9166        let mut bytes = bits.clone();
9167        for g in 0..rows * ng {
9168            let s = 0.02 + 0.001 * g as f32;
9169            bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
9170        }
9171        for r in 0..rows {
9172            let b = bits[r] as usize;
9173            let (mut acc, mut nb) = (0u64, 0usize);
9174            let mut rowbytes = Vec::new();
9175            for i in 0..cols {
9176                let v = ((i * 7 + r * 13) % (1 << b)) as u64;
9177                acc = (acc << b) | v;
9178                nb += b;
9179                while nb >= 8 {
9180                    nb -= 8;
9181                    rowbytes.push(((acc >> nb) & 0xFF) as u8);
9182                }
9183            }
9184            if nb > 0 {
9185                rowbytes.push(((acc << (8 - nb)) & 0xFF) as u8);
9186            }
9187            bytes.extend_from_slice(&rowbytes);
9188        }
9189        let x: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.19).sin()).collect();
9190
9191        let mut reference = vec![0f32; rows * cols];
9192        cortiq_core::quant::dequant_vbit(&bytes, rows, cols, &mut reference).unwrap();
9193        let mut expect = vec![0f32; rows];
9194        for r in 0..rows {
9195            expect[r] = reference[r * cols..(r + 1) * cols]
9196                .iter()
9197                .zip(&x)
9198                .map(|(w, xv)| w * xv)
9199                .sum();
9200        }
9201        let mut got = vec![0f32; rows];
9202        let offsets = vbit_row_offsets(&bytes, rows, cols);
9203        vbitmatvec(&bytes, &offsets, &x, rows, cols, &mut got, None);
9204        // SDOT path quantizes activations to i8 (A8W8): bounded noise,
9205        // same contract as q8 (exact path is pinned by CMF_SDOT=0 in
9206        // the golden-parity gate).
9207        let tol = if a8w8_enabled() { 6e-2 } else { 1e-4 };
9208        let scale = expect.iter().fold(0f32, |m, v| m.max(v.abs())).max(1e-6);
9209        for r in 0..rows {
9210            assert!(
9211                (got[r] - expect[r]).abs() < tol * scale,
9212                "row {r}: {} vs {}",
9213                got[r],
9214                expect[r]
9215            );
9216        }
9217    }
9218
9219    /// Fused q4 matvec must match the reference full-dequant + dense
9220    /// matvec bit-for-bit in structure (same f32 math, group order).
9221    /// vbit matmat: the blocked 1×4 leg must match the per-row path
9222    /// (paired env toggle; larger shape so both code paths engage).
9223    #[test]
9224    #[cfg(target_arch = "x86_64")]
9225    fn vbit_matmat_blocked_matches_per_row() {
9226        let (rows, cols, b) = (64usize, 128usize, 9usize);
9227        let ng = cols / GROUP_SIZE;
9228        let bits: Vec<u8> = (0..rows).map(|r| [3u8, 4, 5, 6][r % 4]).collect();
9229        let mut bytes = bits.clone();
9230        for g in 0..rows * ng {
9231            let sc = 0.02 + 0.0005 * g as f32;
9232            bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(sc).to_le_bytes());
9233        }
9234        for r in 0..rows {
9235            let bw = bits[r] as usize;
9236            let (mut acc, mut nb) = (0u64, 0usize);
9237            let mut rowbytes = Vec::new();
9238            for i in 0..cols {
9239                let v = ((i * 7 + r * 13) % (1 << bw)) as u64;
9240                acc = (acc << bw) | v;
9241                nb += bw;
9242                while nb >= 8 {
9243                    nb -= 8;
9244                    rowbytes.push(((acc >> nb) & 0xFF) as u8);
9245                }
9246            }
9247            if nb > 0 {
9248                rowbytes.push(((acc << (8 - nb)) & 0xFF) as u8);
9249            }
9250            bytes.extend_from_slice(&rowbytes);
9251        }
9252        let x: Vec<f32> = (0..b * cols)
9253            .map(|i| ((i * 13 + 7) % 97) as f32 / 97.0 - 0.5)
9254            .collect();
9255        let offsets = vbit_row_offsets(&bytes, rows, cols);
9256        let mut y_a = vec![0f32; b * rows];
9257        let mut y_b = vec![0f32; b * rows];
9258        unsafe { std::env::set_var("CMF_X86_BLOCKED", "1") };
9259        vbitmatmat(&bytes, &offsets, &x, b, rows, cols, &mut y_a, None);
9260        unsafe { std::env::set_var("CMF_X86_BLOCKED", "0") };
9261        vbitmatmat(&bytes, &offsets, &x, b, rows, cols, &mut y_b, None);
9262        unsafe { std::env::remove_var("CMF_X86_BLOCKED") };
9263        let max_d = y_a
9264            .iter()
9265            .zip(&y_b)
9266            .map(|(p, q)| (p - q).abs())
9267            .fold(0.0f32, f32::max);
9268        assert!(max_d < 1e-4, "vbit blocked ≠ per-row: max|Δ| = {max_d}");
9269    }
9270
9271    /// q4t blocked 1×4 (SDOT on ARM, AVX2 on x86) must equal the
9272    /// per-row path exactly: same nibble unpack, same group order,
9273    /// same f32 accumulation — batch == matvec bit-for-bit. b=9 covers
9274    /// two full 1×4 blocks plus a remainder through the single-row
9275    /// kernel. (Both paths produce identical output, so the shared
9276    /// CMF_X86_BLOCKED env var racing with other tests cannot flip
9277    /// the verdict — worst case both sides take the same path.)
9278    #[test]
9279    fn q4t_matmat_blocked_matches_per_row() {
9280        let (rows, cols, b) = (16usize, 64usize, 9usize);
9281        let gpr = cols / GROUP_SIZE;
9282        let mut bytes = vec![0u8; rows * gpr * Q4_TILE];
9283        for r in 0..rows {
9284            for g in 0..gpr {
9285                let t = (r * gpr + g) * Q4_TILE;
9286                let sc = 0.02 + 0.001 * (r * gpr + g) as f32;
9287                bytes[t..t + 2].copy_from_slice(&cortiq_core::quant::f32_to_f16(sc).to_le_bytes());
9288                for k in 0..16 {
9289                    bytes[t + 2 + k] = ((r * 31 + g * 7 + k * 13) % 251) as u8;
9290                }
9291            }
9292        }
9293        let x: Vec<f32> = (0..b * cols)
9294            .map(|i| ((i * 13 + 7) % 97) as f32 / 97.0 - 0.5)
9295            .collect();
9296        let mut y_blk = vec![0f32; b * rows];
9297        let mut y_row = vec![0f32; b * rows];
9298        unsafe { std::env::set_var("CMF_X86_BLOCKED", "1") };
9299        q4t_matmat(&bytes, &x, b, rows, cols, &mut y_blk, None);
9300        unsafe { std::env::set_var("CMF_X86_BLOCKED", "0") };
9301        q4t_matmat(&bytes, &x, b, rows, cols, &mut y_row, None);
9302        unsafe { std::env::remove_var("CMF_X86_BLOCKED") };
9303        assert_eq!(y_blk, y_row, "q4t blocked 1x4 ≠ per-row");
9304    }
9305
9306    /// The wide-batch Accelerate arm of q4t_matmat vs a brute-force
9307    /// f32 dequant matmul: both are f32 GEMMs, so only reduction
9308    /// order differs — tight tolerance.
9309    /// A synthetic q4tp payload: random nibbles plus a per-row ladder whose
9310    /// span varies row to row, so the codes actually exercise the full 0..31
9311    /// range rather than clustering on one rung.
9312    fn synth_q4tp(rows: usize, cols: usize) -> Vec<u8> {
9313        use cortiq_core::quant::{f32_to_f16, q4tp_code_stride, q4tp_put_code};
9314        let gpr = cols / GROUP_SIZE;
9315        let stride = q4tp_code_stride(gpr);
9316        let (params_off, codes_off, _) = q4tp_sections(rows, cols);
9317        let mut b = vec![0u8; codes_off + rows * stride];
9318        for r in 0..rows {
9319            for g in 0..gpr {
9320                let t = (r * gpr + g) * Q4TP_NIB;
9321                for k in 0..16 {
9322                    b[t + k] = ((r * 31 + g * 7 + k * 13) % 251) as u8;
9323                }
9324            }
9325            let lo = -6.0 - 0.03 * (r % 17) as f32;
9326            let step = 0.01 + 0.004 * (r % 11) as f32;
9327            let p = params_off + r * 4;
9328            b[p..p + 2].copy_from_slice(&f32_to_f16(lo).to_le_bytes());
9329            b[p + 2..p + 4].copy_from_slice(&f32_to_f16(step).to_le_bytes());
9330            let crow = &mut b[codes_off + r * stride..codes_off + (r + 1) * stride];
9331            for g in 0..gpr {
9332                q4tp_put_code(crow, g, (r * 5 + g * 3) % 32);
9333            }
9334        }
9335        b
9336    }
9337
9338    /// The same weights re-expressed as q4_tiled, so the proven kernel can
9339    /// be the reference: each tile stores the ladder scale its code selects.
9340    /// Only the f16 rounding of that scale separates the two payloads.
9341    fn q4tp_as_q4t(bytes: &[u8], rows: usize, cols: usize) -> Vec<u8> {
9342        let gpr = cols / GROUP_SIZE;
9343        let v = Q4tpView::new(bytes, rows, cols);
9344        let mut out = vec![0u8; rows * gpr * Q4_TILE];
9345        let mut sc = vec![0f32; gpr];
9346        for r in 0..rows {
9347            v.scales_into(r, gpr, &mut sc);
9348            for g in 0..gpr {
9349                let t = (r * gpr + g) * Q4_TILE;
9350                let s = sc[g];
9351                out[t..t + 2].copy_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
9352                let src = (r * gpr + g) * Q4TP_NIB;
9353                out[t + 2..t + Q4_TILE].copy_from_slice(&v.nib[src..src + Q4TP_NIB]);
9354            }
9355        }
9356        out
9357    }
9358
9359    /// The exact (`CMF_SDOT=0`) path must reproduce `dequant_q4tp` to f32
9360    /// rounding — that scalar routine is the format's definition, and the
9361    /// kernels re-derive the scale from the ladder independently. Call the
9362    /// row kernel directly: `matmat` picks the int8 arm when a8w8 is on,
9363    /// so routing through it would test the other path by accident.
9364    #[test]
9365    fn q4tp_exact_path_matches_dequant_reference() {
9366        let (rows, cols) = (256usize, 512usize);
9367        let gpr = cols / GROUP_SIZE;
9368        let bytes = synth_q4tp(rows, cols);
9369        let mut w = vec![0f32; rows * cols];
9370        cortiq_core::quant::dequant_q4tp(&bytes, rows, cols, &mut w);
9371
9372        let x: Vec<f32> = (0..cols)
9373            .map(|i| ((i * 13 + 7) % 97) as f32 / 97.0 - 0.5)
9374            .collect();
9375        let v = Q4tpView::new(&bytes, rows, cols);
9376        let mut sc = vec![0f32; gpr];
9377        for r in 0..rows {
9378            v.scales_into(r, gpr, &mut sc);
9379            let got = q4tp_row_exact(v.nib, r, gpr, &x, &sc);
9380            let want: f32 = (0..cols).map(|c| w[r * cols + c] * x[c]).sum();
9381            // These dot products cancel down to ~1e-3 from terms of ~5e-2, so
9382            // the meaningful yardstick is the summed magnitude, not the result:
9383            // against the result any reordering of a 512-term f32 sum "fails".
9384            let mag: f32 = (0..cols).map(|c| (w[r * cols + c] * x[c]).abs()).sum();
9385            assert!(
9386                (got - want).abs() <= 1e-5 * mag,
9387                "row {r}: kernel {got} vs dequant {want}"
9388            );
9389        }
9390    }
9391
9392    /// The int8 (a8w8) path can't be checked against an f32 reference — the
9393    /// activation quantization dominates. Check it against the q4t kernel it
9394    /// was ported from instead, on payloads holding the same weights: that
9395    /// isolates exactly what the port could break (16 B stride, ladder
9396    /// lookup, nibble unpack) from what it deliberately shares.
9397    #[test]
9398    fn q4tp_matvec_matches_the_q4t_kernel_it_was_ported_from() {
9399        let (rows, cols) = (256usize, 512usize);
9400        let bytes = synth_q4tp(rows, cols);
9401        let twin = q4tp_as_q4t(&bytes, rows, cols);
9402        let x: Vec<f32> = (0..cols)
9403            .map(|i| ((i * 13 + 7) % 97) as f32 / 97.0 - 0.5)
9404            .collect();
9405
9406        let mut got = vec![0f32; rows];
9407        q4tp_matvec(&bytes, &x, rows, cols, &mut got, None);
9408        let mut want = vec![0f32; rows];
9409        q4t_matvec(&twin, &x, rows, cols, &mut want, None);
9410
9411        // Scale is f16 in the twin and f32 here, so allow that rounding on
9412        // top of the summed magnitude (same cancellation argument as above).
9413        let mut w = vec![0f32; rows * cols];
9414        cortiq_core::quant::dequant_q4tp(&bytes, rows, cols, &mut w);
9415        for r in 0..rows {
9416            let mag: f32 = (0..cols).map(|c| (w[r * cols + c] * x[c]).abs()).sum();
9417            assert!(
9418                (got[r] - want[r]).abs() <= 1e-3 * mag,
9419                "row {r}: q4tp {} vs q4t {}",
9420                got[r],
9421                want[r]
9422            );
9423        }
9424    }
9425
9426    /// `matmat` carries three arms (Accelerate, blocked int8 1x4, scalar).
9427    /// Batch 5 crosses the blocked kernel's stride, so this exercises the
9428    /// 1x4 path AND its scalar tail in one run — the blocked kernel is new
9429    /// code and its four accumulators are exactly what tends to go wrong.
9430    #[test]
9431    fn q4tp_matmat_matches_the_q4t_kernel_it_was_ported_from() {
9432        let (rows, cols, b) = (256usize, 512usize, 5usize);
9433        let bytes = synth_q4tp(rows, cols);
9434        let twin = q4tp_as_q4t(&bytes, rows, cols);
9435        let xs: Vec<f32> = (0..b * cols)
9436            .map(|i| ((i * 29 + 11) % 89) as f32 / 89.0 - 0.5)
9437            .collect();
9438
9439        let mut got = vec![0f32; b * rows];
9440        q4tp_matmat(&bytes, &xs, b, rows, cols, &mut got, None);
9441        let mut want = vec![0f32; b * rows];
9442        q4t_matmat(&twin, &xs, b, rows, cols, &mut want, None);
9443
9444        let mut w = vec![0f32; rows * cols];
9445        cortiq_core::quant::dequant_q4tp(&bytes, rows, cols, &mut w);
9446        for t in 0..b {
9447            for r in 0..rows {
9448                let mag: f32 = (0..cols)
9449                    .map(|c| (w[r * cols + c] * xs[t * cols + c]).abs())
9450                    .sum();
9451                let (g, wa) = (got[t * rows + r], want[t * rows + r]);
9452                assert!(
9453                    (g - wa).abs() <= 1e-3 * mag,
9454                    "batch {t} row {r}: q4tp {g} vs q4t {wa}"
9455                );
9456            }
9457        }
9458    }
9459
9460    #[test]
9461    fn q4tp_matvec2_matches_the_single_stream_kernel() {
9462        let (rows, cols) = (128usize, 256usize);
9463        let gpr = cols / GROUP_SIZE;
9464        let bytes = synth_q4tp(rows, cols);
9465        let xs: Vec<f32> = (0..2 * cols)
9466            .map(|i| ((i * 29 + 11) % 89) as f32 / 89.0 - 0.5)
9467            .collect();
9468
9469        let (mut o1, mut o2) = (vec![0f32; rows], vec![0f32; rows]);
9470        q4tp_matvec2(
9471            &bytes,
9472            &xs[..cols],
9473            &xs[cols..],
9474            rows,
9475            cols,
9476            &mut o1,
9477            &mut o2,
9478            None,
9479        );
9480
9481        // matvec2 takes the exact path for both streams, so the single-row
9482        // kernel is an exact reference — no tolerance for path differences.
9483        let v = Q4tpView::new(&bytes, rows, cols);
9484        let mut sc = vec![0f32; gpr];
9485        for r in 0..rows {
9486            v.scales_into(r, gpr, &mut sc);
9487            assert_eq!(o1[r], q4tp_row_exact(v.nib, r, gpr, &xs[..cols], &sc));
9488            assert_eq!(o2[r], q4tp_row_exact(v.nib, r, gpr, &xs[cols..], &sc));
9489        }
9490    }
9491
9492    /// q4tp must not COST speed — it exists to save bytes, and a format that
9493    /// trades 7% of a file for a slower model is a bad trade. This guard is
9494    /// here because correctness tests happily passed while `q4tp_matmat` was
9495    /// missing its int8 and Accelerate arms and the model ran 5x slower.
9496    /// Measured on M-series: 0.97-1.04x, i.e. parity (16 B tiles are better
9497    /// aligned than q4t's 18 B, which pays for the scale indirection).
9498    #[test]
9499    fn q4tp_matvec_keeps_pace_with_q4t() {
9500        let (rows, cols) = (4096usize, 3072usize);
9501        let bytes = synth_q4tp(rows, cols);
9502        let twin = q4tp_as_q4t(&bytes, rows, cols);
9503        let x: Vec<f32> = (0..cols).map(|i| (i % 97) as f32 / 97.0 - 0.5).collect();
9504        let mut o = vec![0f32; rows];
9505        let n = 12;
9506        let mut best = (f64::MAX, f64::MAX);
9507        // Interleaved A/B, minimum statistic: this machine throttles, and a
9508        // mean over a thermal ramp reliably indicts whichever ran second.
9509        for _ in 0..3 {
9510            let t0 = std::time::Instant::now();
9511            for _ in 0..n {
9512                q4t_matvec(&twin, &x, rows, cols, &mut o, None);
9513            }
9514            best.0 = best.0.min(t0.elapsed().as_secs_f64());
9515            let t0 = std::time::Instant::now();
9516            for _ in 0..n {
9517                q4tp_matvec(&bytes, &x, rows, cols, &mut o, None);
9518            }
9519            best.1 = best.1.min(t0.elapsed().as_secs_f64());
9520        }
9521        let ratio = best.1 / best.0;
9522        println!(
9523            "q4t {:.3} ms | q4tp {:.3} ms | {ratio:.2}x",
9524            best.0 * 1e3 / n as f64,
9525            best.1 * 1e3 / n as f64
9526        );
9527        assert!(ratio < 2.0, "q4tp matvec {ratio:.2}x slower than q4t");
9528    }
9529
9530    #[cfg(target_os = "macos")]
9531    #[test]
9532    fn q4t_matmat_accel_matches_dequant_reference() {
9533        if !accel_gemm_enabled() {
9534            return; // CMF_ACCEL=0
9535        }
9536        let (rows, cols, b) = (512usize, 1024usize, 8usize); // ≥500K → accel arm
9537        let gpr = cols / GROUP_SIZE;
9538        let mut bytes = vec![0u8; rows * gpr * Q4_TILE];
9539        for r in 0..rows {
9540            for g in 0..gpr {
9541                let t = (r * gpr + g) * Q4_TILE;
9542                let sc = 0.02 + 0.0005 * ((r * gpr + g) % 64) as f32;
9543                bytes[t..t + 2].copy_from_slice(&cortiq_core::quant::f32_to_f16(sc).to_le_bytes());
9544                for k in 0..16 {
9545                    bytes[t + 2 + k] = ((r * 31 + g * 7 + k * 13) % 251) as u8;
9546                }
9547            }
9548        }
9549        let x: Vec<f32> = (0..b * cols)
9550            .map(|i| ((i * 13 + 7) % 97) as f32 / 97.0 - 0.5)
9551            .collect();
9552        let mut got = vec![0f32; b * rows];
9553        q4t_matmat(&bytes, &x, b, rows, cols, &mut got, None);
9554        // Brute-force reference off the same tiles.
9555        let mut w = vec![0f32; rows * cols];
9556        for r in 0..rows {
9557            for g in 0..gpr {
9558                let t = (r * gpr + g) * Q4_TILE;
9559                let s = f16_to_f32(u16::from_le_bytes([bytes[t], bytes[t + 1]]));
9560                for (k, &bb) in bytes[t + 2..t + Q4_TILE].iter().enumerate() {
9561                    w[r * cols + g * GROUP_SIZE + k * 2] = ((bb & 0x0F) as f32 - 8.0) * s;
9562                    w[r * cols + g * GROUP_SIZE + k * 2 + 1] =
9563                        (((bb >> 4) & 0x0F) as f32 - 8.0) * s;
9564                }
9565            }
9566        }
9567        for bi in 0..b {
9568            for r in 0..rows {
9569                let want: f32 = (0..cols).map(|j| x[bi * cols + j] * w[r * cols + j]).sum();
9570                let d = (got[bi * rows + r] - want).abs();
9571                assert!(
9572                    d <= want.abs().max(1.0) * 1e-4,
9573                    "accel q4t GEMM diverged at ({bi},{r}): {} vs {want}",
9574                    got[bi * rows + r]
9575                );
9576            }
9577        }
9578    }
9579
9580    #[test]
9581    fn q4matvec_matches_full_dequant() {
9582        let (rows, cols) = (8, 64);
9583        let groups = rows * cols / GROUP_SIZE;
9584        // Hand-craft a q4_block blob: nibbles then f16 scales.
9585        let mut bytes = Vec::with_capacity(groups * 16 + groups * 2);
9586        for i in 0..groups * 16 {
9587            bytes.push((((i * 7 + 3) % 256) & 0xFF) as u8);
9588        }
9589        for g in 0..groups {
9590            let s = 0.01 + 0.003 * g as f32;
9591            bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
9592        }
9593        let x: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.17).sin()).collect();
9594
9595        let mut reference = vec![0.0f32; rows * cols];
9596        cortiq_core::quant::dequant_q4_block(&bytes, &mut reference);
9597        let mut expect = vec![0.0f32; rows];
9598        for r in 0..rows {
9599            expect[r] = reference[r * cols..(r + 1) * cols]
9600                .iter()
9601                .zip(&x)
9602                .map(|(w, xv)| w * xv)
9603                .sum();
9604        }
9605
9606        let mut got = vec![0.0f32; rows];
9607        q4matvec(&bytes, &x, rows, cols, &mut got, None);
9608        // SDOT path quantizes activations to i8 (A8W8): bounded noise,
9609        // same contract as q8/vbit (exact path is pinned by CMF_SDOT=0
9610        // in the golden-parity gate).
9611        let tol = if a8w8_enabled() { 6e-2 } else { 1e-4 };
9612        let scale = expect.iter().fold(0f32, |m, v| m.max(v.abs())).max(1.0);
9613        for r in 0..rows {
9614            assert!(
9615                (got[r] - expect[r]).abs() < tol * scale,
9616                "row {r}: {} vs {}",
9617                got[r],
9618                expect[r]
9619            );
9620        }
9621    }
9622
9623    /// Fused two-input vbit matvec must equal two single matvecs exactly
9624    /// (same per-lane accumulation order on both scalar and SDOT paths).
9625    #[test]
9626    fn vbitmatvec2_equals_two_singles() {
9627        let (rows, cols) = (6, 64);
9628        let ng = cols / GROUP_SIZE;
9629        let bits: Vec<u8> = vec![3, 4, 5, 6, 8, 4];
9630        let mut bytes = bits.clone();
9631        for g in 0..rows * ng {
9632            let s = 0.02 + 0.001 * g as f32;
9633            bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
9634        }
9635        for r in 0..rows {
9636            let b = bits[r] as usize;
9637            let (mut acc, mut nb) = (0u64, 0usize);
9638            let mut rowbytes = Vec::new();
9639            for i in 0..cols {
9640                let v = ((i * 7 + r * 13) % (1 << b)) as u64;
9641                acc = (acc << b) | v;
9642                nb += b;
9643                while nb >= 8 {
9644                    nb -= 8;
9645                    rowbytes.push(((acc >> nb) & 0xFF) as u8);
9646                }
9647            }
9648            if nb > 0 {
9649                rowbytes.push(((acc << (8 - nb)) & 0xFF) as u8);
9650            }
9651            bytes.extend_from_slice(&rowbytes);
9652        }
9653        let x1: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.19).sin()).collect();
9654        let x2: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.11).cos()).collect();
9655        let offsets = vbit_row_offsets(&bytes, rows, cols);
9656
9657        let (mut a1, mut a2) = (vec![0f32; rows], vec![0f32; rows]);
9658        vbitmatvec(&bytes, &offsets, &x1, rows, cols, &mut a1, None);
9659        vbitmatvec(&bytes, &offsets, &x2, rows, cols, &mut a2, None);
9660        let (mut b1, mut b2) = (vec![0f32; rows], vec![0f32; rows]);
9661        vbitmatvec2(
9662            &bytes, &offsets, &x1, &x2, rows, cols, &mut b1, &mut b2, None,
9663        );
9664        assert_eq!(a1, b1, "fused vbit lane 1 must be bit-identical");
9665        assert_eq!(a2, b2, "fused vbit lane 2 must be bit-identical");
9666    }
9667
9668    /// Fused two-input q4 matvec must equal two single matvecs exactly.
9669    #[test]
9670    fn q4matvec2_equals_two_singles() {
9671        let (rows, cols) = (8, 128);
9672        let groups = rows * cols / GROUP_SIZE;
9673        let mut bytes = Vec::with_capacity(groups * 16 + groups * 2);
9674        for i in 0..groups * 16 {
9675            bytes.push((((i * 7 + 3) % 256) & 0xFF) as u8);
9676        }
9677        for g in 0..groups {
9678            let s = 0.01 + 0.003 * g as f32;
9679            bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
9680        }
9681        // Include an outlier channel so the SDOT correction path is
9682        // exercised in the pair kernel too.
9683        let mut x1: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.17).sin()).collect();
9684        x1[9] = 250.0;
9685        let x2: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.23).cos()).collect();
9686
9687        let (mut a1, mut a2) = (vec![0f32; rows], vec![0f32; rows]);
9688        q4matvec(&bytes, &x1, rows, cols, &mut a1, None);
9689        q4matvec(&bytes, &x2, rows, cols, &mut a2, None);
9690        let (mut b1, mut b2) = (vec![0f32; rows], vec![0f32; rows]);
9691        q4matvec2(&bytes, &x1, &x2, rows, cols, &mut b1, &mut b2, None);
9692        assert_eq!(a1, b1, "fused q4 lane 1 must be bit-identical");
9693        assert_eq!(a2, b2, "fused q4 lane 2 must be bit-identical");
9694    }
9695
9696    /// Multi-matrix job must equal separate matvecs exactly — same
9697    /// kernels, only the dispatch is fused.
9698    #[test]
9699    fn matvec_many_equals_separate_matvecs() {
9700        use crate::pool::Pool;
9701        let (r1, r2, cols) = (300, 200, 64);
9702        let mk = |salt: usize, rows: usize| {
9703            QTensor::from_f32(
9704                (0..rows * cols)
9705                    .map(|i| ((i * 7 + salt) % 97) as f32 / 97.0 - 0.5)
9706                    .collect(),
9707                rows,
9708                cols,
9709            )
9710        };
9711        let (a, b) = (mk(1, r1), mk(5, r2));
9712        let x: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.11).sin()).collect();
9713        let pool = Pool::new(3);
9714
9715        let (mut ea, mut eb) = (vec![0f32; r1], vec![0f32; r2]);
9716        a.matvec(&x, &mut ea, Some(&pool));
9717        b.matvec(&x, &mut eb, Some(&pool));
9718        let (mut ga, mut gb) = (vec![0f32; r1], vec![0f32; r2]);
9719        QTensor::matvec_many([&a, &b], &x, [&mut ga, &mut gb], Some(&pool));
9720        assert_eq!(ea, ga, "fused multi-matrix lane 1 must be bit-identical");
9721        assert_eq!(eb, gb, "fused multi-matrix lane 2 must be bit-identical");
9722    }
9723
9724    /// Batched q4/vbit matmat must equal per-position matvec calls
9725    /// exactly (the fallback it replaced) — same kernels, same order.
9726    #[test]
9727    fn batched_matmat_equals_per_position_matvec() {
9728        let (rows, cols, b) = (8, 64, 5);
9729        // q4 blob.
9730        let groups = rows * cols / GROUP_SIZE;
9731        let mut q4 = Vec::new();
9732        for i in 0..groups * 16 {
9733            q4.push((((i * 7 + 3) % 256) & 0xFF) as u8);
9734        }
9735        for g in 0..groups {
9736            q4.extend_from_slice(
9737                &cortiq_core::quant::f32_to_f16(0.01 + 0.003 * g as f32).to_le_bytes(),
9738            );
9739        }
9740        // vbit blob (mixed widths incl. 8).
9741        let ng = cols / GROUP_SIZE;
9742        let bits: Vec<u8> = vec![3, 4, 5, 6, 8, 4, 5, 3];
9743        let mut vb = bits.clone();
9744        for g in 0..rows * ng {
9745            vb.extend_from_slice(
9746                &cortiq_core::quant::f32_to_f16(0.02 + 0.001 * g as f32).to_le_bytes(),
9747            );
9748        }
9749        for r in 0..rows {
9750            let bw = bits[r] as usize;
9751            let (mut acc, mut nb) = (0u64, 0usize);
9752            let mut rowbytes = Vec::new();
9753            for i in 0..cols {
9754                let v = ((i * 7 + r * 13) % (1 << bw)) as u64;
9755                acc = (acc << bw) | v;
9756                nb += bw;
9757                while nb >= 8 {
9758                    nb -= 8;
9759                    rowbytes.push(((acc >> nb) & 0xFF) as u8);
9760                }
9761            }
9762            if nb > 0 {
9763                rowbytes.push(((acc << (8 - nb)) & 0xFF) as u8);
9764            }
9765            vb.extend_from_slice(&rowbytes);
9766        }
9767        let offsets = vbit_row_offsets(&vb, rows, cols);
9768
9769        let xs: Vec<f32> = (0..b * cols).map(|i| (i as f32 * 0.13).sin()).collect();
9770
9771        // q4: batch vs singles.
9772        let mut got = vec![0f32; b * rows];
9773        q4matmat(&q4, &xs, b, rows, cols, &mut got, None);
9774        for bi in 0..b {
9775            let mut expect = vec![0f32; rows];
9776            q4matvec(
9777                &q4,
9778                &xs[bi * cols..(bi + 1) * cols],
9779                rows,
9780                cols,
9781                &mut expect,
9782                None,
9783            );
9784            assert_eq!(
9785                &got[bi * rows..(bi + 1) * rows],
9786                &expect[..],
9787                "q4 batch pos {bi}"
9788            );
9789        }
9790
9791        // vbit: batch vs singles.
9792        let mut got = vec![0f32; b * rows];
9793        vbitmatmat(&vb, &offsets, &xs, b, rows, cols, &mut got, None);
9794        for bi in 0..b {
9795            let mut expect = vec![0f32; rows];
9796            vbitmatvec(
9797                &vb,
9798                &offsets,
9799                &xs[bi * cols..(bi + 1) * cols],
9800                rows,
9801                cols,
9802                &mut expect,
9803                None,
9804            );
9805            assert_eq!(
9806                &got[bi * rows..(bi + 1) * rows],
9807                &expect[..],
9808                "vbit batch pos {bi}"
9809            );
9810        }
9811    }
9812
9813    /// q4_tiled kernels must produce BIT-identical outputs to the q4
9814    /// split kernels on the same values (same ints, same order — only
9815    /// the byte placement differs).
9816    #[test]
9817    fn q4_tiled_matches_q4_block_bitexact() {
9818        let (rows, cols, b) = (8usize, 128usize, 3usize);
9819        let groups = rows * cols / GROUP_SIZE;
9820        let mut split = Vec::with_capacity(groups * 18);
9821        for i in 0..groups * 16 {
9822            split.push((((i * 7 + 3) % 256) & 0xFF) as u8);
9823        }
9824        for g in 0..groups {
9825            split.extend_from_slice(
9826                &cortiq_core::quant::f32_to_f16(0.01 + 0.003 * g as f32).to_le_bytes(),
9827            );
9828        }
9829        // Re-tile: [scale][nibbles] per group.
9830        let (packed, scales) = split.split_at(groups * 16);
9831        let mut tiled = Vec::with_capacity(groups * Q4_TILE);
9832        for g in 0..groups {
9833            tiled.extend_from_slice(&scales[g * 2..g * 2 + 2]);
9834            tiled.extend_from_slice(&packed[g * 16..(g + 1) * 16]);
9835        }
9836
9837        let mut x1: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.17).sin()).collect();
9838        x1[9] = 250.0; // exercise the outlier path
9839        let x2: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.23).cos()).collect();
9840
9841        let (mut a, mut t) = (vec![0f32; rows], vec![0f32; rows]);
9842        q4matvec(&split, &x1, rows, cols, &mut a, None);
9843        q4t_matvec(&tiled, &x1, rows, cols, &mut t, None);
9844        assert_eq!(a, t, "q4t matvec must match q4 bit-for-bit");
9845
9846        let (mut a1, mut a2) = (vec![0f32; rows], vec![0f32; rows]);
9847        let (mut t1, mut t2) = (vec![0f32; rows], vec![0f32; rows]);
9848        q4matvec2(&split, &x1, &x2, rows, cols, &mut a1, &mut a2, None);
9849        q4t_matvec2(&tiled, &x1, &x2, rows, cols, &mut t1, &mut t2, None);
9850        assert_eq!(a1, t1);
9851        assert_eq!(a2, t2);
9852
9853        let xs: Vec<f32> = (0..b * cols).map(|i| (i as f32 * 0.13).sin()).collect();
9854        let (mut am, mut tm) = (vec![0f32; b * rows], vec![0f32; b * rows]);
9855        q4matmat(&split, &xs, b, rows, cols, &mut am, None);
9856        q4t_matmat(&tiled, &xs, b, rows, cols, &mut tm, None);
9857        assert_eq!(am, tm, "q4t matmat must match q4 bit-for-bit");
9858    }
9859
9860    /// q4 SDOT outlier correction: a single huge activation channel
9861    /// (>8·rms → outlier, zeroed in xq) must still contribute its EXACT
9862    /// term. On-grid bulk (±1/0 → xq dequantizes exactly) isolates the
9863    /// correction from A8W8 noise. cols must exceed 64: at n=64 the
9864    /// 8·rms threshold equals sqrt(v²+rest) ≥ v, so a single outlier
9865    /// can never qualify (8² = n).
9866    #[test]
9867    fn q4matvec_sdot_outlier_exact() {
9868        let (rows, cols) = (4, 128);
9869        let groups = rows * cols / GROUP_SIZE;
9870        let mut bytes = Vec::with_capacity(groups * 16 + groups * 2);
9871        for i in 0..groups * 16 {
9872            bytes.push(((i * 11 + 5) % 256) as u8);
9873        }
9874        for g in 0..groups {
9875            let s = 0.02 + 0.002 * g as f32;
9876            bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
9877        }
9878        let mut x: Vec<f32> = (0..cols)
9879            .map(|i| match i % 3 {
9880                0 => 1.0,
9881                1 => -1.0,
9882                _ => 0.0,
9883            })
9884            .collect();
9885        x[17] = 300.0; // ≫ 8·rms → outlier channel
9886
9887        let mut reference = vec![0.0f32; rows * cols];
9888        cortiq_core::quant::dequant_q4_block(&bytes, &mut reference);
9889        let mut expect = vec![0.0f32; rows];
9890        for r in 0..rows {
9891            expect[r] = reference[r * cols..(r + 1) * cols]
9892                .iter()
9893                .zip(&x)
9894                .map(|(w, xv)| w * xv)
9895                .sum();
9896        }
9897        let mut got = vec![0.0f32; rows];
9898        q4matvec(&bytes, &x, rows, cols, &mut got, None);
9899        let scale = expect.iter().fold(0f32, |m, v| m.max(v.abs())).max(1.0);
9900        for r in 0..rows {
9901            assert!(
9902                (got[r] - expect[r]).abs() < 2e-3 * scale,
9903                "row {r}: {} vs {} (outlier term must be exact)",
9904                got[r],
9905                expect[r]
9906            );
9907        }
9908    }
9909
9910    /// The fused q1t matvec must equal the reference (dequant_q1t → dot),
9911    /// including the ternary zero level and the binary-searched outlier
9912    /// overlay. Guards the mmap kernel that makes a 12B q1t runnable.
9913    #[test]
9914    fn q1t_matvec_matches_reference() {
9915        use cortiq_core::quant::{dequant_q1t, f32_to_f16};
9916        let (rows, cols) = (3usize, 64usize); // gpr = 2
9917        let gpr = cols / GROUP_SIZE;
9918        let scales = [0.5f32, 0.3, 0.7, 0.2, 0.6, 0.15];
9919        // Overlay (must be sorted by flat index): a few spikes across rows.
9920        let outliers: [(u32, f32); 3] = [(5, 9.0), (70, -4.5), (150, 3.25)];
9921        let is_out = |flat: usize| outliers.iter().any(|&(i, _)| i as usize == flat);
9922        let mut bytes = Vec::new();
9923        for r in 0..rows {
9924            for g in 0..gpr {
9925                bytes.extend_from_slice(&f32_to_f16(scales[r * gpr + g]).to_le_bytes());
9926                let mut c = [0u8; 7];
9927                for k in 0..GROUP_SIZE {
9928                    // Encoder invariant: code 0 at outlier positions.
9929                    let code = if is_out(r * cols + g * GROUP_SIZE + k) {
9930                        0
9931                    } else {
9932                        ((k + r * 3 + g) % 3) as u8 // 0,1,2
9933                    };
9934                    cortiq_core::quant::q1t_pack(&mut c, k, code);
9935                }
9936                bytes.extend_from_slice(&c);
9937            }
9938        }
9939        // Per-row overlay: [u32 row_ptr[rows+1]] then [(u16 col, f16 val)] by
9940        // row (outliers are sorted by flat index → already grouped by row).
9941        let mut row_ptr = vec![0u32; rows + 1];
9942        for &(idx, _) in &outliers {
9943            row_ptr[idx as usize / cols + 1] += 1;
9944        }
9945        for r in 0..rows {
9946            row_ptr[r + 1] += row_ptr[r];
9947        }
9948        for &p in &row_ptr {
9949            bytes.extend_from_slice(&p.to_le_bytes());
9950        }
9951        for &(idx, v) in &outliers {
9952            bytes.extend_from_slice(&((idx as usize % cols) as u16).to_le_bytes());
9953            bytes.extend_from_slice(&f32_to_f16(v).to_le_bytes());
9954        }
9955
9956        let mut refw = vec![0f32; rows * cols];
9957        dequant_q1t(&bytes, rows, cols, &mut refw);
9958        // On-grid activations (±1, amax 1) so the int8 SDOT path reconstructs
9959        // x exactly and matches the f32 reference (same trick as the q1 test).
9960        let x: Vec<f32> = (0..cols)
9961            .map(|j| if j % 3 == 0 { 1.0 } else { -1.0 })
9962            .collect();
9963        let mut expect = vec![0f32; rows];
9964        for r in 0..rows {
9965            let mut a = 0.0f32;
9966            for j in 0..cols {
9967                a += refw[r * cols + j] * x[j];
9968            }
9969            expect[r] = a;
9970        }
9971        let tol = |e: f32| 1e-3 * e.abs().max(1e-3);
9972        let mut got = vec![0f32; rows];
9973        q1t_matvec(&bytes, &x, rows, cols, &mut got, None);
9974        for r in 0..rows {
9975            assert!(
9976                (got[r] - expect[r]).abs() < tol(expect[r]),
9977                "row {r}: {} vs {}",
9978                got[r],
9979                expect[r]
9980            );
9981        }
9982        // matmat (b=2, f32 decode path) must agree too.
9983        let x2: Vec<f32> = x.iter().chain(x.iter().map(|v| v)).copied().collect();
9984        let mut gm = vec![0f32; 2 * rows];
9985        q1t_matmat(&bytes, &x2, 2, rows, cols, &mut gm, None);
9986        for r in 0..rows {
9987            assert!((gm[r] - expect[r]).abs() < tol(expect[r]));
9988            assert!((gm[rows + r] - expect[r]).abs() < tol(expect[r]));
9989        }
9990        // Fused pair (q1t_matvec2) must equal two single matvecs
9991        // bit-for-bit: same unpack, same group order, same f32
9992        // accumulation per stream. Distinct x2 exercises both lanes.
9993        let xb: Vec<f32> = (0..cols)
9994            .map(|j| if j % 5 == 0 { -1.0 } else { 1.0 })
9995            .collect();
9996        let (mut s1, mut s2) = (vec![0f32; rows], vec![0f32; rows]);
9997        q1t_matvec(&bytes, &x, rows, cols, &mut s1, None);
9998        q1t_matvec(&bytes, &xb, rows, cols, &mut s2, None);
9999        let (mut p1, mut p2) = (vec![0f32; rows], vec![0f32; rows]);
10000        q1t_matvec2(&bytes, &x, &xb, rows, cols, &mut p1, &mut p2, None);
10001        assert_eq!(p1, s1, "q1t pair lane 1 ≠ single matvec");
10002        assert_eq!(p2, s2, "q1t pair lane 2 ≠ single matvec");
10003    }
10004
10005    /// Pair == 2×matvec with an ODD group count (the kernel's tail
10006    /// group) and no overlay section.
10007    #[test]
10008    fn q1t_matvec2_odd_gpr_matches_singles() {
10009        use cortiq_core::quant::{Q1T_TILE, f32_to_f16, q1t_pack};
10010        let (rows, cols) = (5usize, 96usize); // gpr = 3 → paired + tail
10011        let gpr = cols / GROUP_SIZE;
10012        let mut bytes = Vec::with_capacity(rows * gpr * Q1T_TILE);
10013        for r in 0..rows {
10014            for g in 0..gpr {
10015                bytes.extend_from_slice(&f32_to_f16(0.1 + 0.05 * (r + g) as f32).to_le_bytes());
10016                let mut c = [0u8; 7];
10017                for k in 0..GROUP_SIZE {
10018                    q1t_pack(&mut c, k, ((k * 7 + r * 5 + g * 3) % 3) as u8);
10019                }
10020                bytes.extend_from_slice(&c);
10021            }
10022        }
10023        let x1: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.31).sin()).collect();
10024        let x2: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.17).cos()).collect();
10025        let (mut s1, mut s2) = (vec![0f32; rows], vec![0f32; rows]);
10026        q1t_matvec(&bytes, &x1, rows, cols, &mut s1, None);
10027        q1t_matvec(&bytes, &x2, rows, cols, &mut s2, None);
10028        let (mut p1, mut p2) = (vec![0f32; rows], vec![0f32; rows]);
10029        q1t_matvec2(&bytes, &x1, &x2, rows, cols, &mut p1, &mut p2, None);
10030        assert_eq!(p1, s1, "odd-gpr pair lane 1 ≠ single");
10031        assert_eq!(p2, s2, "odd-gpr pair lane 2 ≠ single");
10032    }
10033
10034    // Speed A/B: fused pair (one unpack, two streams) vs two single
10035    // matvecs. Single-threaded, FFN-sized, min-of paired in-process.
10036    //   cargo test -p cortiq-engine --release q1t_matvec2_speed -- --ignored --nocapture
10037    #[test]
10038    #[ignore]
10039    fn q1t_matvec2_speed() {
10040        use cortiq_core::quant::{Q1T_TILE, f32_to_f16, q1t_pack};
10041        use std::time::Instant;
10042        let (rows, cols) = (8192usize, 4096usize);
10043        let gpr = cols / GROUP_SIZE;
10044        let mut bytes = Vec::with_capacity(rows * gpr * Q1T_TILE);
10045        for r in 0..rows {
10046            for g in 0..gpr {
10047                let s = 0.1 + ((r + g) % 7) as f32 * 0.01;
10048                bytes.extend_from_slice(&f32_to_f16(s).to_le_bytes());
10049                let mut c = [0u8; 7];
10050                for k in 0..GROUP_SIZE {
10051                    q1t_pack(&mut c, k, ((k * 7 + r + g) % 3) as u8);
10052                }
10053                bytes.extend_from_slice(&c);
10054            }
10055        }
10056        let x1: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.31).sin()).collect();
10057        let x2: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.17).cos()).collect();
10058        let (mut s1, mut s2) = (vec![0f32; rows], vec![0f32; rows]);
10059        let (mut p1, mut p2) = (vec![0f32; rows], vec![0f32; rows]);
10060        // Warm both paths once.
10061        q1t_matvec(&bytes, &x1, rows, cols, &mut s1, None);
10062        q1t_matvec2(&bytes, &x1, &x2, rows, cols, &mut p1, &mut p2, None);
10063        let (mut t_pair, mut t_two) = (f64::MAX, f64::MAX);
10064        for _ in 0..8 {
10065            let t0 = Instant::now();
10066            q1t_matvec2(&bytes, &x1, &x2, rows, cols, &mut p1, &mut p2, None);
10067            t_pair = t_pair.min(t0.elapsed().as_secs_f64() * 1000.0);
10068            let t1 = Instant::now();
10069            q1t_matvec(&bytes, &x1, rows, cols, &mut s1, None);
10070            q1t_matvec(&bytes, &x2, rows, cols, &mut s2, None);
10071            t_two = t_two.min(t1.elapsed().as_secs_f64() * 1000.0);
10072        }
10073        assert_eq!(p1, s1);
10074        assert_eq!(p2, s2);
10075        println!("q1t pair {rows}x{cols}: fused {t_pair:.2} ms | two singles {t_two:.2} ms");
10076    }
10077
10078    // Speed A/B: the base-3-division decode (what the packing commit left in
10079    // place) vs the fused sign-LUT matvec. Both single-threaded, same bytes.
10080    //   cargo test -p cortiq-engine q1t_matvec_speed -- --ignored --nocapture
10081    #[test]
10082    #[ignore]
10083    fn q1t_matvec_speed() {
10084        use cortiq_core::quant::{Q1T_TILE, f32_to_f16, q1t_code, q1t_pack};
10085        use std::time::Instant;
10086        let (rows, cols) = (8192usize, 4096usize); // FFN-sized
10087        let gpr = cols / GROUP_SIZE;
10088        let mut bytes = Vec::with_capacity(rows * gpr * Q1T_TILE + 16);
10089        for r in 0..rows {
10090            for g in 0..gpr {
10091                let s = 0.1 + ((r + g) % 7) as f32 * 0.01;
10092                bytes.extend_from_slice(&f32_to_f16(s).to_le_bytes());
10093                let mut c = [0u8; 7];
10094                for k in 0..GROUP_SIZE {
10095                    q1t_pack(&mut c, k, ((k * 7 + r + g) % 3) as u8);
10096                }
10097                bytes.extend_from_slice(&c);
10098            }
10099        }
10100        let (n, stride) = (rows * cols, 40usize); // ~2.5% outliers, per-row overlay
10101        let mut row_ptr = vec![0u32; rows + 1];
10102        let mut idx = 0usize;
10103        while idx < n {
10104            row_ptr[idx / cols + 1] += 1;
10105            idx += stride;
10106        }
10107        for r in 0..rows {
10108            row_ptr[r + 1] += row_ptr[r];
10109        }
10110        for &p in &row_ptr {
10111            bytes.extend_from_slice(&p.to_le_bytes());
10112        }
10113        let mut idx = 0usize;
10114        while idx < n {
10115            bytes.extend_from_slice(&((idx % cols) as u16).to_le_bytes());
10116            bytes.extend_from_slice(&f32_to_f16((idx % 13) as f32 * 0.1 - 0.6).to_le_bytes());
10117            idx += stride;
10118        }
10119        // On-grid ±1 so the fast path's int8 SDOT is exact vs the f32 "slow"
10120        // reference (the A/B is a timing check; values must still agree).
10121        let x: Vec<f32> = (0..cols)
10122            .map(|j| if j % 3 == 0 { 1.0 } else { -1.0 })
10123            .collect();
10124        let (rp_off, ent_off, has_ov) = q1t_overlay(&bytes, rows * gpr * Q1T_TILE, rows);
10125
10126        // "before": base-3 division decode into a buffer, then dot.
10127        let slow = |out: &mut [f32]| {
10128            let mut buf = vec![0f32; cols];
10129            for r in 0..rows {
10130                for g in 0..gpr {
10131                    let off = (r * gpr + g) * Q1T_TILE;
10132                    let s = f16_to_f32(u16::from_le_bytes([bytes[off], bytes[off + 1]]));
10133                    let codes = &bytes[off + 2..off + Q1T_TILE];
10134                    for k in 0..GROUP_SIZE {
10135                        buf[g * GROUP_SIZE + k] = match q1t_code(codes, k) {
10136                            1 => s,
10137                            2 => -s,
10138                            _ => 0.0,
10139                        };
10140                    }
10141                }
10142                out[r] = q1t_row_outlier_correction(&bytes, r, rp_off, ent_off, has_ov, &x)
10143                    + (0..cols).map(|j| buf[j] * x[j]).sum::<f32>();
10144            }
10145        };
10146        let iters = 5;
10147        let mut a = vec![0f32; rows];
10148        slow(&mut a); // warm
10149        let t = Instant::now();
10150        for _ in 0..iters {
10151            slow(&mut a);
10152        }
10153        let slow_ms = t.elapsed().as_secs_f64() * 1e3 / iters as f64;
10154
10155        let mut b = vec![0f32; rows];
10156        q1t_matvec(&bytes, &x, rows, cols, &mut b, None); // warm
10157        let t = Instant::now();
10158        for _ in 0..iters {
10159            q1t_matvec(&bytes, &x, rows, cols, &mut b, None);
10160        }
10161        let fast_ms = t.elapsed().as_secs_f64() * 1e3 / iters as f64;
10162
10163        for r in 0..rows {
10164            assert!((a[r] - b[r]).abs() < 1e-2, "mismatch row {r}");
10165        }
10166        println!(
10167            "q1t matvec {rows}x{cols} (1 thread): div-decode {slow_ms:.2} ms  fused-LUT {fast_ms:.2} ms  => {:.2}x",
10168            slow_ms / fast_ms
10169        );
10170    }
10171}