Skip to main content

cortiq_engine/
qtensor.rs

1//! QTensor — weight tensor with pluggable storage.
2//!
3//! Two backings, one interface:
4//! - `F32`   — owned dense floats (small models, tests). Every operation
5//!   is bit-identical to the historical `&[f32]` code paths.
6//! - `Mapped` — quantized bytes zero-copy from the CMF mmap (`q8_row` /
7//!   `q8_2f`). The matvec is fused: int8 rows × f32 activations, the
8//!   q8_2f column field folds into a pre-scale of the input
9//!   (`x'[i] = col[i]·x[i]`), so the inner loop is the same i8 dot as
10//!   q8_row. This is what lets a 15B file run in a few GB of RSS.
11//!
12//! Extension point: new dtypes = new match arm here, nothing else moves.
13
14use crate::pool::{Pool, matvec_rows, matvec_rows2};
15use cortiq_core::quant::{
16    GROUP_SIZE, Q1_TILE, Q2TP_CHUNK, Q4_TILE, Q4TP_NIB, f16_to_f32, q2tp_ladder, q2tp_sections,
17    q4tp_code, q4tp_ladder, q4tp_sections,
18};
19use cortiq_core::{CmfModel, TensorDtype};
20use std::sync::Arc;
21
22pub enum QTensor {
23    F32 {
24        data: Vec<f32>,
25        rows: usize,
26        cols: usize,
27    },
28    Mapped {
29        model: Arc<CmfModel>,
30        /// Index into the model's tensor directory.
31        idx: usize,
32        dtype: TensorDtype,
33        rows: usize,
34        cols: usize,
35        /// Per-row scales, dequantized to f32 up front (tiny).
36        row_scale: Vec<f32>,
37        /// q8_2f column field (θ), dequantized up front; empty for q8_row.
38        col_field: Vec<f32>,
39        /// Vbit only: byte offset of each row's packed data within the
40        /// tensor blob (`[rows + 1]`, computed once at load — the per-
41        /// matvec prefix scan over row bit-widths was O(rows) each call).
42        vbit_offsets: Vec<usize>,
43        /// q8-family decode repack (load-time, optional): rows in groups
44        /// of 4, interleaved in 16-byte units — one 64-byte line per
45        /// iteration feeds all 4 sdot lanes, ONE sequential weight
46        /// stream per worker instead of four (this is where llama.cpp's
47        /// repacked Q8 kernels get their bandwidth). Empty = off
48        /// (CMF_REPACK=0, non-SDOT arch, or an ineligible shape). Trades
49        /// an anonymous copy of the quants for mmap pages that go cold.
50        repack: Vec<u8>,
51    },
52}
53
54/// Load-time q8 repack gate (see `Mapped::repack`). OPT-IN
55/// (`CMF_REPACK=1`): the single-stream hypothesis LOST on Apple Silicon
56/// (M4, interleaved A/B: decode 101 vs 94 tok/s — four adjacent row
57/// streams per worker feed the prefetcher MORE memory-level parallelism
58/// than one); kept as an experiment flag for x86, where the tradeoff
59/// may land differently.
60fn repack_enabled() -> bool {
61    static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
62    *ON.get_or_init(|| {
63        std::env::var("CMF_REPACK")
64            .map(|v| v == "1")
65            .unwrap_or(cfg!(target_os = "android"))
66    })
67}
68
69/// Interleave q8 rows for the decode kernel: group g holds rows
70/// 4g..4g+4 as [r0[c], r1[c], r2[c], r3[c]] per 16-byte chunk c. Only
71/// full groups are packed — tail rows keep reading the mmap layout.
72fn q8_repack(bytes: &[u8], rows: usize, cols: usize) -> Vec<u8> {
73    #[cfg(target_arch = "aarch64")]
74    let arch_ok = sdot_enabled();
75    #[cfg(not(target_arch = "aarch64"))]
76    let arch_ok = false;
77    if !arch_ok || !repack_enabled() || rows < 256 || cols % 16 != 0 {
78        return Vec::new();
79    }
80    q8_repack_layout(bytes, rows, cols)
81}
82
83/// The pure layout transform behind `q8_repack` (tested directly —
84/// the gate depends on arch and env).
85fn q8_repack_layout(bytes: &[u8], rows: usize, cols: usize) -> Vec<u8> {
86    let groups = rows / 4;
87    let mut rep = vec![0u8; groups * 4 * cols];
88    for g in 0..groups {
89        let dst = &mut rep[g * 4 * cols..(g + 1) * 4 * cols];
90        for c in 0..cols / 16 {
91            for lane in 0..4 {
92                let src = (g * 4 + lane) * cols + c * 16;
93                dst[c * 64 + lane * 16..c * 64 + lane * 16 + 16]
94                    .copy_from_slice(&bytes[src..src + 16]);
95            }
96        }
97    }
98    rep
99}
100
101/// Prefix-sum of vbit row payload offsets (absolute within the tensor
102/// bytes). `offsets[r]..offsets[r+1]` is row r's packed data.
103fn vbit_row_offsets(bytes: &[u8], rows: usize, cols: usize) -> Vec<usize> {
104    let ng = cols / GROUP_SIZE;
105    let bits = &bytes[..rows];
106    let mut offsets = Vec::with_capacity(rows + 1);
107    let mut off = rows + rows * ng * 2;
108    for r in 0..rows {
109        offsets.push(off);
110        off += (cols * bits[r] as usize).div_ceil(8);
111    }
112    offsets.push(off);
113    offsets
114}
115
116/// `CMF_X86_BLOCKED` / `CMF_GPU_LMHEAD` / `CMF_GPU_SPLIT`, read once. They
117/// used to be read from the environment on every large matvec and on every
118/// matmat in six places — microseconds each, but also a knob that could
119/// change under a running process, which is not a thing a kernel choice
120/// should be able to do mid-sequence.
121fn blocked_enabled() -> bool {
122    use std::sync::atomic::Ordering::Relaxed;
123    match BLOCKED_OVERRIDE.load(Relaxed) {
124        1 => false,
125        2 => true,
126        _ => {
127            static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
128            *ON.get_or_init(|| {
129                std::env::var("CMF_X86_BLOCKED")
130                    .map(|v| v != "0")
131                    .unwrap_or(true)
132            })
133        }
134    }
135}
136
137static BLOCKED_OVERRIDE: std::sync::atomic::AtomicU8 = std::sync::atomic::AtomicU8::new(0);
138
139/// Force the blocked GEMM on or off, ignoring the environment; `None`
140/// restores it. For tests that need to run BOTH paths and compare them:
141/// `blocked_enabled` caches its answer for the life of the process, which
142/// is right when the environment is the only input, but leaves a test that
143/// flips `CMF_X86_BLOCKED` between two calls comparing a path against
144/// itself — or against whatever a test running in parallel latched first.
145pub fn set_blocked_override(on: Option<bool>) {
146    let v = match on {
147        None => 0,
148        Some(false) => 1,
149        Some(true) => 2,
150    };
151    BLOCKED_OVERRIDE.store(v, std::sync::atomic::Ordering::Relaxed);
152}
153
154fn gpu_lmhead_enabled() -> bool {
155    static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
156    *ON.get_or_init(|| {
157        std::env::var("CMF_GPU_LMHEAD")
158            .map(|v| v != "0")
159            .unwrap_or(true)
160    })
161}
162
163fn gpu_split_frac() -> f32 {
164    static F: std::sync::OnceLock<f32> = std::sync::OnceLock::new();
165    *F.get_or_init(|| {
166        std::env::var("CMF_GPU_SPLIT")
167            .ok()
168            .and_then(|v| v.parse::<f32>().ok())
169            .unwrap_or(0.5)
170            .clamp(0.0, 1.0)
171    })
172}
173
174impl QTensor {
175    pub fn from_f32(data: Vec<f32>, rows: usize, cols: usize) -> Self {
176        debug_assert_eq!(data.len(), rows * cols);
177        Self::F32 { data, rows, cols }
178    }
179
180    /// Wrap a directory tensor without dequantizing the payload.
181    /// Falls back to dequantized f32 for dtypes without a fused kernel.
182    pub fn from_model(model: &Arc<CmfModel>, name: &str) -> Result<Self, String> {
183        // Indexed lookup: the linear directory scan made pipeline build
184        // O(N²) on MoE/skills files with thousands of tensors.
185        let idx = model
186            .tensor_index(name)
187            .ok_or_else(|| format!("tensor '{name}' not found in CMF directory"))?;
188        let entry = &model.tensors[idx];
189        if entry.shape.len() != 2 {
190            return Err(format!("QTensor::from_model needs 2-D, got '{name}'"));
191        }
192        let (rows, cols) = (entry.shape[0], entry.shape[1]);
193        let bytes = model.entry_bytes(entry);
194
195        match entry.dtype {
196            TensorDtype::Q8Row | TensorDtype::Q8_2f => {
197                let n = rows * cols;
198                let scales_off = n;
199                let row_scale: Vec<f32> = (0..rows)
200                    .map(|o| {
201                        f16_to_f32(u16::from_le_bytes([
202                            bytes[scales_off + o * 2],
203                            bytes[scales_off + o * 2 + 1],
204                        ]))
205                    })
206                    .collect();
207                let col_field: Vec<f32> = if entry.dtype == TensorDtype::Q8_2f {
208                    let col_off = n + rows * 2;
209                    (0..cols)
210                        .map(|i| {
211                            f16_to_f32(u16::from_le_bytes([
212                                bytes[col_off + i * 2],
213                                bytes[col_off + i * 2 + 1],
214                            ]))
215                        })
216                        .collect()
217                } else {
218                    Vec::new()
219                };
220                Ok(Self::Mapped {
221                    model: model.clone(),
222                    idx,
223                    dtype: entry.dtype,
224                    rows,
225                    cols,
226                    row_scale,
227                    col_field,
228                    vbit_offsets: Vec::new(),
229                    repack: q8_repack(bytes, rows, cols),
230                })
231            }
232            // vbit: fused kernel unpacks variable-bit rows from mmap.
233            TensorDtype::Vbit if cols % GROUP_SIZE == 0 => Ok(Self::Mapped {
234                model: model.clone(),
235                idx,
236                dtype: entry.dtype,
237                rows,
238                cols,
239                row_scale: Vec::new(),
240                col_field: Vec::new(),
241                vbit_offsets: vbit_row_offsets(bytes, rows, cols),
242                repack: Vec::new(),
243            }),
244            // vbit_ro (§4.2): the offset table comes straight from the
245            // file — no load-time prefix scan; kernels are shared with
246            // legacy vbit (they consume absolute offsets either way).
247            TensorDtype::VbitRo if cols % GROUP_SIZE == 0 => {
248                let (_, off_off, packed_off) = cortiq_core::quant::vbit_ro_sections(rows, cols);
249                let offsets: Vec<usize> = (0..=rows)
250                    .map(|r| packed_off + cortiq_core::quant::vbit_ro_offset(bytes, off_off, r))
251                    .collect();
252                Ok(Self::Mapped {
253                    model: model.clone(),
254                    idx,
255                    dtype: entry.dtype,
256                    rows,
257                    cols,
258                    row_scale: Vec::new(),
259                    col_field: Vec::new(),
260                    vbit_offsets: offsets,
261                    repack: Vec::new(),
262                })
263            }
264            // q4_block: fused kernel reads nibbles straight from mmap —
265            // a 14B q4 file no longer explodes into ×8 f32 RAM.
266            // q4_tiled (§4.3): interleaved [scale][nibbles] tiles — one
267            // sequential memory stream (measured ×1.66 ARM / ×1.13 AVX2
268            // at kernel level over the split layout).
269            TensorDtype::Q4Tiled if cols % GROUP_SIZE == 0 => Ok(Self::Mapped {
270                model: model.clone(),
271                idx,
272                dtype: entry.dtype,
273                rows,
274                cols,
275                row_scale: Vec::new(),
276                col_field: Vec::new(),
277                vbit_offsets: Vec::new(),
278                repack: Vec::new(),
279            }),
280            // q4tp (§4.10): nibbles from mmap, scale from the row ladder —
281            // 7.3% less file than q4t at the same 4-bit grid.
282            TensorDtype::Q4TiledP if cols % GROUP_SIZE == 0 => Ok(Self::Mapped {
283                model: model.clone(),
284                idx,
285                dtype: entry.dtype,
286                rows,
287                cols,
288                row_scale: Vec::new(),
289                col_field: Vec::new(),
290                vbit_offsets: Vec::new(),
291                repack: Vec::new(),
292            }),
293            // q2tp: 2-bit chunks from mmap, scale from the same row ladder.
294            TensorDtype::Q2TiledP if cols % GROUP_SIZE == 0 => Ok(Self::Mapped {
295                model: model.clone(),
296                idx,
297                dtype: entry.dtype,
298                rows,
299                cols,
300                row_scale: Vec::new(),
301                col_field: Vec::new(),
302                vbit_offsets: Vec::new(),
303                repack: Vec::new(),
304            }),
305            TensorDtype::Q4Block if cols % GROUP_SIZE == 0 => Ok(Self::Mapped {
306                model: model.clone(),
307                idx,
308                dtype: entry.dtype,
309                rows,
310                cols,
311                row_scale: Vec::new(),
312                col_field: Vec::new(),
313                vbit_offsets: Vec::new(),
314                repack: Vec::new(),
315            }),
316            // q1: binary sign-bit tiles from mmap (1-bit-trained models).
317            TensorDtype::Q1 if cols % GROUP_SIZE == 0 => Ok(Self::Mapped {
318                model: model.clone(),
319                idx,
320                dtype: entry.dtype,
321                rows,
322                cols,
323                row_scale: Vec::new(),
324                col_field: Vec::new(),
325                vbit_offsets: Vec::new(),
326                repack: Vec::new(),
327            }),
328            // q1t (ternary + outlier overlay): fused per-row dequant kernel
329            // reads straight from mmap — a 12B q1t stays ~its file size in
330            // RAM instead of dequantizing to ~48 GB of f32.
331            TensorDtype::Q1T if cols % GROUP_SIZE == 0 => Ok(Self::Mapped {
332                model: model.clone(),
333                idx,
334                dtype: entry.dtype,
335                rows,
336                cols,
337                row_scale: Vec::new(),
338                col_field: Vec::new(),
339                vbit_offsets: Vec::new(),
340                repack: Vec::new(),
341            }),
342            // No fused kernel yet → dequantize once (correct, more RAM).
343            _ => {
344                let mut data = vec![0.0f32; rows * cols];
345                cortiq_core::quant::dequant_tensor(entry, bytes, &mut data)?;
346                Ok(Self::from_f32(data, rows, cols))
347            }
348        }
349    }
350
351    /// q1-mapped tensor? (GPU gates: the q1 CPU kernel is
352    /// compute-bound, so offload pays at much smaller shapes than q8.)
353    pub(crate) fn is_q1(&self) -> bool {
354        matches!(
355            self,
356            Self::Mapped {
357                dtype: TensorDtype::Q1,
358                ..
359            }
360        )
361    }
362
363    /// Owned-f32 view (data, rows, cols) — the GDN a/b gate projections
364    /// arrive dequantized (force-f16 in the converter → F32 in RAM).
365    pub(crate) fn f32_parts(&self) -> Option<(&[f32], usize, usize)> {
366        match self {
367            Self::F32 { data, rows, cols } => Some((data, *rows, *cols)),
368            _ => None,
369        }
370    }
371
372    /// (directory idx, rows, cols) of a q1-mapped tensor — the
373    /// whole-block GPU path resolves offsets itself.
374    /// (idx, rows, cols) of a mapped tensor the whole-token GPU graph can drive
375    /// — Q1, Q1T or Q4-block (it resolves the offset and picks the kernel by
376    /// dtype). Q4-block lets a precise down_proj/lm_head stay on-device.
377    /// Named `q1_parts` for historical reasons.
378    pub(crate) fn q1_parts(&self) -> Option<(usize, usize, usize)> {
379        match self {
380            #[cfg(target_os = "macos")]
381            Self::Mapped {
382                dtype: TensorDtype::Q1T,
383                ..
384            } if !crate::gpu::metal_q1t_enabled() => None,
385            Self::Mapped {
386                idx,
387                dtype:
388                    TensorDtype::Q1
389                    | TensorDtype::Q1T
390                    | TensorDtype::Q4Block
391                    | TensorDtype::Q4Tiled
392                    // Q2TiledP deliberately absent: the Metal graph has no
393                    // q2tp kernel, and advertising it here made the block
394                    // plan truncate mid-run at the first q2tp layer.
395                    | TensorDtype::Q4TiledP
396                    | TensorDtype::Q8Row
397                    | TensorDtype::Q8_2f,
398                rows,
399                cols,
400                ..
401            } => Some((*idx, *rows, *cols)),
402            _ => None,
403        }
404    }
405
406    /// (directory idx, rows, cols) of a q4_tiled mapped tensor. The
407    /// chunk-prefill graph takes it in the same 4-tuple slot as
408    /// `q8_row_parts` with an EMPTY row_scale — q4t carries its scales
409    /// inside the 18-byte tiles, and the empty slice is what tells the
410    /// encoder to reach for the q4t kernels.
411    pub(crate) fn q4t_parts(&self) -> Option<(usize, usize, usize)> {
412        match self {
413            Self::Mapped {
414                idx,
415                dtype: TensorDtype::Q4Tiled,
416                rows,
417                cols,
418                ..
419            } => Some((*idx, *rows, *cols)),
420            _ => None,
421        }
422    }
423
424    /// (directory idx, rows, cols) of a q4tp mapped tensor. Same empty-scale
425    /// slot as `q4t_parts` in the chunk graph — the encoder tells the two
426    /// apart by the tensor's dtype, not by the slot.
427    pub(crate) fn q4tp_parts(&self) -> Option<(usize, usize, usize)> {
428        match self {
429            Self::Mapped {
430                idx,
431                dtype: TensorDtype::Q4TiledP,
432                rows,
433                cols,
434                ..
435            } => Some((*idx, *rows, *cols)),
436            _ => None,
437        }
438    }
439
440    /// (directory idx, rows, cols, row_scale) of a plain q8_row mapped
441    /// tensor — the chunk-prefill GPU graph resolves offsets itself.
442    /// q8_2f is excluded on purpose: its column field would need a
443    /// prescale stage on the device.
444    pub(crate) fn q8_row_parts(&self) -> Option<(usize, usize, usize, &[f32])> {
445        match self {
446            Self::Mapped {
447                idx,
448                dtype: TensorDtype::Q8Row,
449                rows,
450                cols,
451                row_scale,
452                col_field,
453                ..
454            } if col_field.is_empty() => Some((*idx, *rows, *cols, row_scale)),
455            _ => None,
456        }
457    }
458
459    /// The layout this tensor is stored in, when it is mapped from a model.
460    /// The frames branch on it — a q2tp gate against a q4tp down is a real
461    /// combination in the 2-bit profile and needs a different kernel.
462    pub fn model_dtype(&self) -> Option<cortiq_core::TensorDtype> {
463        match self {
464            Self::Mapped { dtype, .. } => Some(*dtype),
465            _ => None,
466        }
467    }
468
469    /// The tensor's index in the model directory, when it is mapped from one.
470    /// The GPU frames bind by index rather than by name — a name lookup per
471    /// layer per token is not free, and the index is what the device cache is
472    /// keyed on anyway.
473    pub fn model_idx(&self) -> Option<usize> {
474        match self {
475            Self::Mapped { idx, .. } => Some(*idx),
476            _ => None,
477        }
478    }
479
480    /// The model this tensor is mapped from, when it is mapped at all. The
481    /// GPU frames need the container to reach the bytes; a QTensor already
482    /// holds it, and threading a second handle down every call site to say
483    /// the same thing invites the two to disagree.
484    pub fn model_arc(&self) -> Option<std::sync::Arc<cortiq_core::CmfModel>> {
485        match self {
486            Self::Mapped { model, .. } => Some(model.clone()),
487            _ => None,
488        }
489    }
490
491    pub fn rows(&self) -> usize {
492        match self {
493            Self::F32 { rows, .. } | Self::Mapped { rows, .. } => *rows,
494        }
495    }
496
497    /// Mapped q4t handle (model + directory index) — the fused GPU FFN
498    /// needs the raw file coordinates of its three projections.
499    pub(crate) fn mapped_q4t(&self) -> Option<(&Arc<CmfModel>, usize)> {
500        match self {
501            Self::Mapped {
502                model,
503                idx,
504                dtype: TensorDtype::Q4Tiled,
505                ..
506            } => Some((model, *idx)),
507            _ => None,
508        }
509    }
510
511    /// Same slot as `mapped_q4t` for a q4tp tensor — the fused DiT FFN picks
512    /// its kernels by which of the two answers.
513    pub fn mapped_q4tp(&self) -> Option<(&Arc<CmfModel>, usize)> {
514        match self {
515            Self::Mapped {
516                model,
517                idx,
518                dtype: TensorDtype::Q4TiledP,
519                ..
520            } => Some((model, *idx)),
521            _ => None,
522        }
523    }
524
525    /// (model, tensor idx) for a q2tp mapped weight — the 2-bit twin of
526    /// `mapped_q4tp`, used by the mixed MoE profile.
527    pub fn mapped_q2tp(&self) -> Option<(&Arc<CmfModel>, usize)> {
528        match self {
529            Self::Mapped {
530                model,
531                idx,
532                dtype: TensorDtype::Q2TiledP,
533                ..
534            } => Some((model, *idx)),
535            _ => None,
536        }
537    }
538
539    pub fn cols(&self) -> usize {
540        match self {
541            Self::F32 { cols, .. } | Self::Mapped { cols, .. } => *cols,
542        }
543    }
544
545    /// (model, tensor idx) for a q1 mapped weight — the wgpu token graph
546    /// keys its resident VRAM cache by idx. None for any other dtype/kind.
547    pub fn mapped_q1(&self) -> Option<(&std::sync::Arc<CmfModel>, usize)> {
548        match self {
549            Self::Mapped {
550                model,
551                idx,
552                dtype: TensorDtype::Q1,
553                ..
554            } => Some((model, *idx)),
555            _ => None,
556        }
557    }
558
559    /// (model, idx, kind, row_scale) for a graph-capable mapped weight.
560    /// kind: 0=q8_row (per-row scales), 1=q1, 2=q4_block, 3=q1t
561    /// (tile-embedded, no rs), 5=q4_tiled, 6=q4tp, 7=q8_2f (both scale
562    /// planes live inside the tensor). None only for `vbit`.
563    ///
564    /// The old comment here claimed q4_block was unhandled while the arm
565    /// right below mapped it, and it named q8_2f as unhandled after that
566    /// stopped being true — a stale comment on this function is how a
567    /// model silently loses the graph, so it is worth keeping honest.
568    pub fn graph_weight(&self) -> Option<(&std::sync::Arc<CmfModel>, usize, u8, &[f32])> {
569        match self {
570            Self::Mapped {
571                model,
572                idx,
573                dtype: TensorDtype::Q8Row,
574                row_scale,
575                ..
576            } => Some((model, *idx, 0, row_scale.as_slice())),
577            Self::Mapped {
578                model,
579                idx,
580                dtype: TensorDtype::Q1,
581                ..
582            } => Some((model, *idx, 1, &[])),
583            // Q4Tiled is kind 5, NOT 2: both carried 2 historically, and
584            // the wgpu token graph fed 18B interleaved tiles to the
585            // split-layout q4b kernel — garbage output on q4t models
586            // (caught by an end-to-end answer check on real Vulkan).
587            Self::Mapped {
588                model,
589                idx,
590                dtype: TensorDtype::Q4Tiled,
591                ..
592            } => Some((model, *idx, 5, &[])),
593            // Kind 6, not 5: q4tp's nibble stride and scale planes differ,
594            // and feeding them to the q4t kernel is exactly the mistake that
595            // produced garbage when Q4Tiled shared kind 2 with Q4Block.
596            Self::Mapped {
597                model,
598                idx,
599                dtype: TensorDtype::Q4TiledP,
600                ..
601            } => Some((model, *idx, 6, &[])),
602            Self::Mapped {
603                model,
604                idx,
605                dtype: TensorDtype::Q4Block,
606                ..
607            } => Some((model, *idx, 2, &[])),
608            // q8_2f carries BOTH scale planes after the int8 body (rows
609            // f16, then cols f16), so the graph takes the whole tensor
610            // and the kernel reads them where they lie — no host-side
611            // prescale, which is what the per-op path does instead.
612            Self::Mapped {
613                model,
614                idx,
615                dtype: TensorDtype::Q8_2f,
616                ..
617            } => Some((model, *idx, 7, &[])),
618            Self::Mapped {
619                model,
620                idx,
621                dtype: TensorDtype::Q1T,
622                ..
623            } => Some((model, *idx, 3, &[])),
624            // Kind 9: the 2-bit plane on the q4tp ladder (dense FFN gate/up
625            // of the q2tp profile). Its own kernel — 8 bytes a group where
626            // q4tp has 16, and rung 0 is the exact zero.
627            Self::Mapped {
628                model,
629                idx,
630                dtype: TensorDtype::Q2TiledP,
631                ..
632            } => Some((model, *idx, 9, &[])),
633            _ => None,
634        }
635    }
636
637    /// Dense f32 view — only for owned tensors. Masked/sparse execution
638    /// paths require it; quantized weights don't support masks yet.
639    pub fn as_f32(&self) -> Option<&[f32]> {
640        match self {
641            Self::F32 { data, .. } => Some(data),
642            Self::Mapped { .. } => None,
643        }
644    }
645
646    fn quant_bytes(&self) -> &[u8] {
647        match self {
648            Self::Mapped { model, idx, .. } => model.entry_bytes(&model.tensors[*idx]),
649            Self::F32 { .. } => unreachable!("quant_bytes on F32"),
650        }
651    }
652
653    /// Dequantize one row into `dst` (embedding lookup).
654    pub fn row_f32(&self, r: usize, dst: &mut [f32]) {
655        let cols = self.cols();
656        debug_assert_eq!(dst.len(), cols);
657        match self {
658            Self::F32 { data, .. } => dst.copy_from_slice(&data[r * cols..(r + 1) * cols]),
659            Self::Mapped {
660                dtype,
661                row_scale,
662                col_field,
663                vbit_offsets,
664                ..
665            } => {
666                if *dtype == TensorDtype::Q4Tiled {
667                    let bytes = self.quant_bytes();
668                    let gpr = cols / GROUP_SIZE;
669                    for gi in 0..gpr {
670                        let tile = &bytes[(r * gpr + gi) * Q4_TILE..(r * gpr + gi + 1) * Q4_TILE];
671                        let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
672                        for (k, &b) in tile[2..].iter().enumerate() {
673                            dst[gi * GROUP_SIZE + k * 2] = ((b & 0x0F) as f32 - 8.0) * s;
674                            dst[gi * GROUP_SIZE + k * 2 + 1] = (((b >> 4) & 0x0F) as f32 - 8.0) * s;
675                        }
676                    }
677                    return;
678                }
679                if *dtype == TensorDtype::Q4TiledP {
680                    let bytes = self.quant_bytes();
681                    let gpr = cols / GROUP_SIZE;
682                    let v = Q4tpView::new(bytes, self.rows(), cols);
683                    let mut sc = vec![0f32; gpr];
684                    v.scales_into(r, gpr, &mut sc);
685                    for gi in 0..gpr {
686                        let tile = &v.nib[(r * gpr + gi) * Q4TP_NIB..(r * gpr + gi + 1) * Q4TP_NIB];
687                        let s = sc[gi];
688                        for (k, &b) in tile.iter().enumerate() {
689                            dst[gi * GROUP_SIZE + k * 2] = ((b & 0x0F) as f32 - 8.0) * s;
690                            dst[gi * GROUP_SIZE + k * 2 + 1] = (((b >> 4) & 0x0F) as f32 - 8.0) * s;
691                        }
692                    }
693                    return;
694                }
695                if *dtype == TensorDtype::Q2TiledP {
696                    let bytes = self.quant_bytes();
697                    let gpr = cols / GROUP_SIZE;
698                    let v = Q4tpView::new_q2(bytes, self.rows(), cols);
699                    let mut sc = vec![0f32; gpr];
700                    v.scales_into(r, gpr, &mut sc);
701                    for gi in 0..gpr {
702                        let ch =
703                            &v.nib[(r * gpr + gi) * Q2TP_CHUNK..(r * gpr + gi + 1) * Q2TP_CHUNK];
704                        let s = sc[gi];
705                        for (k, &b) in ch.iter().enumerate() {
706                            for j in 0..4 {
707                                dst[gi * GROUP_SIZE + k * 4 + j] =
708                                    (((b >> (2 * j)) & 3) as f32 - 1.5) * s;
709                            }
710                        }
711                    }
712                    return;
713                }
714                if *dtype == TensorDtype::Q4Block {
715                    let (packed, scales) = q4_split(self.quant_bytes(), self.rows(), cols);
716                    let gpr = cols / GROUP_SIZE;
717                    for gi in 0..gpr {
718                        let g = r * gpr + gi;
719                        let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
720                        for (k, &b) in packed[g * 16..(g + 1) * 16].iter().enumerate() {
721                            dst[gi * GROUP_SIZE + k * 2] = ((b & 0x0F) as f32 - 8.0) * s;
722                            dst[gi * GROUP_SIZE + k * 2 + 1] = (((b >> 4) & 0x0F) as f32 - 8.0) * s;
723                        }
724                    }
725                    return;
726                }
727                if *dtype == TensorDtype::Q1 {
728                    let bytes = self.quant_bytes();
729                    let gpr = cols / GROUP_SIZE;
730                    for gi in 0..gpr {
731                        let tile = &bytes[(r * gpr + gi) * Q1_TILE..(r * gpr + gi + 1) * Q1_TILE];
732                        let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
733                        for (j, &b) in tile[2..].iter().enumerate() {
734                            for k in 0..8 {
735                                dst[gi * GROUP_SIZE + j * 8 + k] =
736                                    (((b >> k) & 1) as f32 * 2.0 - 1.0) * s;
737                            }
738                        }
739                    }
740                    return;
741                }
742                if *dtype == TensorDtype::Q1T {
743                    let bytes = self.quant_bytes();
744                    let gpr = cols / GROUP_SIZE;
745                    let base_len = self.rows() * gpr * cortiq_core::quant::Q1T_TILE;
746                    for gi in 0..gpr {
747                        let off = (r * gpr + gi) * cortiq_core::quant::Q1T_TILE;
748                        let s = cortiq_core::quant::f16_to_f32(u16::from_le_bytes([
749                            bytes[off],
750                            bytes[off + 1],
751                        ]));
752                        let codes = &bytes[off + 2..off + cortiq_core::quant::Q1T_TILE];
753                        for k in 0..GROUP_SIZE {
754                            dst[gi * GROUP_SIZE + k] = match cortiq_core::quant::q1t_code(codes, k)
755                            {
756                                1 => s,
757                                2 => -s,
758                                _ => 0.0,
759                            };
760                        }
761                    }
762                    // Overlay
763                    let rows = self.rows();
764                    let entries = base_len + (rows + 1) * 4;
765                    if entries <= bytes.len() {
766                        let ptrs = &bytes[base_len..base_len + (rows + 1) * 4];
767                        let r0 = u32::from_le_bytes([
768                            ptrs[r * 4],
769                            ptrs[r * 4 + 1],
770                            ptrs[r * 4 + 2],
771                            ptrs[r * 4 + 3],
772                        ]) as usize;
773                        let r1 = u32::from_le_bytes([
774                            ptrs[(r + 1) * 4],
775                            ptrs[(r + 1) * 4 + 1],
776                            ptrs[(r + 1) * 4 + 2],
777                            ptrs[(r + 1) * 4 + 3],
778                        ]) as usize;
779                        let off = entries + r0 * 4;
780                        for i in 0..r1 - r0 {
781                            let item = &bytes[off + i * 4..off + i * 4 + 4];
782                            let c = u16::from_le_bytes([item[0], item[1]]) as usize;
783                            let v = cortiq_core::quant::f16_to_f32(u16::from_le_bytes([
784                                item[2], item[3],
785                            ]));
786                            if c < cols {
787                                dst[c] = v;
788                            }
789                        }
790                    }
791                    return;
792                }
793                if matches!(dtype, TensorDtype::Vbit | TensorDtype::VbitRo) {
794                    let bytes = self.quant_bytes();
795                    let rows = self.rows();
796                    let ng = cols / GROUP_SIZE;
797                    let bits = &bytes[..rows];
798                    let sc_off = rows;
799                    // Precomputed at load — embedding lookup used to scan
800                    // the bit-widths of every preceding row (O(token_id)).
801                    let off = vbit_offsets[r];
802                    let b = bits[r] as usize;
803                    let l = ((1usize << (b - 1)) - 1) as f32;
804                    let data = &bytes[off..];
805                    let (mut acc, mut nbits, mut idx) = (0u64, 0usize, 0usize);
806                    for (i, d) in dst.iter_mut().enumerate() {
807                        while nbits < b {
808                            acc = (acc << 8) | data[idx] as u64;
809                            idx += 1;
810                            nbits += 8;
811                        }
812                        let u = ((acc >> (nbits - b)) & ((1u64 << b) - 1)) as f32;
813                        nbits -= b;
814                        let so = (r * ng + i / GROUP_SIZE) * 2;
815                        let sv = f16_to_f32(u16::from_le_bytes([
816                            bytes[sc_off + so],
817                            bytes[sc_off + so + 1],
818                        ]));
819                        *d = (u - l) * sv;
820                    }
821                    return;
822                }
823                let q = &self.quant_bytes()[r * cols..(r + 1) * cols];
824                let s = row_scale[r];
825                match dtype {
826                    TensorDtype::Q8Row => {
827                        for (d, &b) in dst.iter_mut().zip(q) {
828                            *d = (b as i8) as f32 * s;
829                        }
830                    }
831                    TensorDtype::Q8_2f => {
832                        for (i, (d, &b)) in dst.iter_mut().zip(q).enumerate() {
833                            *d = (b as i8) as f32 * s * col_field[i];
834                        }
835                    }
836                    _ => unreachable!(),
837                }
838            }
839        }
840    }
841
842    /// Can this tensor's columns be read cheaply (for sparse down_proj)?
843    /// True for F32/Q8Row/Q8_2f (per-row scale, direct strided access);
844    /// false for group-packed q4/vbit (column access would unpack whole
845    /// groups — sparse execution falls back to f32 for those).
846    pub fn sparse_col_ok(&self) -> bool {
847        match self {
848            Self::F32 { .. } => true,
849            Self::Mapped { dtype, .. } => {
850                matches!(dtype, TensorDtype::Q8Row | TensorDtype::Q8_2f)
851            }
852        }
853    }
854
855    /// down_proj [hidden, inter]: accumulate `w · col(c)` into `out`
856    /// [hidden] — reads ONLY column `c` (one neuron) from the mmap,
857    /// no full-matrix dequant. `out[k] += w · down[k, c]`.
858    pub fn add_col_scaled(&self, c: usize, w: f32, out: &mut [f32]) {
859        let inter = self.cols();
860        let hidden = self.rows();
861        debug_assert_eq!(out.len(), hidden);
862        match self {
863            Self::F32 { data, .. } => {
864                for (k, o) in out.iter_mut().enumerate() {
865                    *o += w * data[k * inter + c];
866                }
867            }
868            Self::Mapped {
869                dtype,
870                row_scale,
871                col_field,
872                ..
873            } => {
874                let q = self.quant_bytes();
875                let colf = if *dtype == TensorDtype::Q8_2f {
876                    col_field[c]
877                } else {
878                    1.0
879                };
880                let wc = w * colf;
881                for (k, o) in out.iter_mut().enumerate() {
882                    let b = q[k * inter + c] as i8 as f32;
883                    *o += wc * b * row_scale[k];
884                }
885            }
886        }
887    }
888
889    /// Dot of row `r` with `x` (gate/up active-neuron path). Reads only
890    /// row `r` from the mmap — no full dequant. q4/vbit dequant the row
891    /// into `scratch` first (rare for active-FFN weights).
892    pub fn row_dot(&self, r: usize, x: &[f32], scratch: &mut [f32]) -> f32 {
893        let cols = self.cols();
894        match self {
895            Self::F32 { data, .. } => {
896                let row = &data[r * cols..(r + 1) * cols];
897                row.iter().zip(x).map(|(w, v)| w * v).sum()
898            }
899            Self::Mapped {
900                dtype,
901                row_scale,
902                col_field,
903                ..
904            } => match dtype {
905                TensorDtype::Q8Row => {
906                    let q = &self.quant_bytes()[r * cols..(r + 1) * cols];
907                    dot_i8_f32(q, x) * row_scale[r]
908                }
909                TensorDtype::Q8_2f => {
910                    let q = &self.quant_bytes()[r * cols..(r + 1) * cols];
911                    dot_i8_col_f32(q, x, col_field) * row_scale[r]
912                }
913                _ => {
914                    self.row_f32(r, scratch);
915                    scratch.iter().zip(x).map(|(w, v)| w * v).sum()
916                }
917            },
918        }
919    }
920
921    /// `out = W · x` (row-major). F32 delegates to the historical
922    /// bit-exact path; Mapped runs the fused int8 kernel.
923    pub fn matvec(&self, x: &[f32], out: &mut [f32], pool: Option<&Pool>) {
924        match self {
925            // NOTE: `out.len()` DRIVES this arm — it computes that many rows,
926            // and `x.len()` is the stride. A short `out` is legitimate here,
927            // which is why the check below lives in the Mapped arm only.
928            Self::F32 { data, .. } => matvec_rows(pool, data, x, out),
929            Self::Mapped {
930                model,
931                idx,
932                dtype,
933                rows,
934                cols,
935                row_scale,
936                col_field,
937                vbit_offsets,
938                repack,
939            } => {
940                let _ = (model, idx);
941                // Every kernel below writes `rows` entries through a raw
942                // pointer, so a short `out` is an out-of-bounds WRITE, not a
943                // wrong answer: it scribbles on the allocator's metadata and
944                // the process aborts much later, somewhere innocent
945                // (`double free or corruption`, `corrupted double-linked
946                // list`). The debug_assert two of the kernels carried is
947                // compiled out of the release — exactly the build where it
948                // matters. Fail here instead, while the caller is still on
949                // the stack to be named.
950                assert!(
951                    out.len() >= *rows && x.len() >= *cols,
952                    "matvec {rows}x{cols}: out {} (need {rows}), x {} (need {cols})",
953                    out.len(),
954                    x.len(),
955                );
956                if *dtype == TensorDtype::Q4Block {
957                    // GPU route (wgpu q4b kernel) for large q4_block matvecs —
958                    // gives NVIDIA/AMD/Intel q4 models a GPU path. Probe keeps
959                    // the winner; Metal returns false → the CPU kernel below.
960                    if *rows * *cols >= 8_388_608 && crate::gpu::enabled_here() {
961                        let t0 = std::time::Instant::now();
962                        match crate::gpu::probe_arm(crate::gpu::OpClass::Matvec) {
963                            crate::gpu::ProbeArm::Gpu => {
964                                if crate::gpu::q4b_matvec(model, *idx, x, *rows, *cols, out) {
965                                    crate::gpu::probe_record(
966                                        crate::gpu::OpClass::Matvec,
967                                        true,
968                                        t0.elapsed(),
969                                    );
970                                    return;
971                                }
972                            }
973                            crate::gpu::ProbeArm::CpuTimed => {
974                                q4matvec(self.quant_bytes(), x, *rows, *cols, out, pool);
975                                crate::gpu::probe_record(
976                                    crate::gpu::OpClass::Matvec,
977                                    false,
978                                    t0.elapsed(),
979                                );
980                                return;
981                            }
982                            crate::gpu::ProbeArm::Cpu => {}
983                        }
984                    }
985                    q4matvec(self.quant_bytes(), x, *rows, *cols, out, pool);
986                    return;
987                }
988                if *dtype == TensorDtype::Q4Tiled {
989                    // GPU route for large q4t matvecs — the lm_head class,
990                    // same shape as the q4tp arm below. The probe keeps the
991                    // winner; a backend without the kernel refuses and the
992                    // CPU path stays.
993                    if *rows * *cols >= 8_388_608 && crate::gpu::enabled_here() {
994                        let t0 = std::time::Instant::now();
995                        let cls = crate::gpu::matvec_class(*rows, *cols);
996                        match crate::gpu::probe_arm(cls) {
997                            crate::gpu::ProbeArm::Gpu => {
998                                if crate::gpu::q4t_matvec(model, *idx, x, *rows, *cols, out) {
999                                    crate::gpu::probe_record(cls, true, t0.elapsed());
1000                                    return;
1001                                }
1002                            }
1003                            crate::gpu::ProbeArm::CpuTimed => {
1004                                q4t_matvec(self.quant_bytes(), x, *rows, *cols, out, pool);
1005                                crate::gpu::probe_record(cls, false, t0.elapsed());
1006                                return;
1007                            }
1008                            crate::gpu::ProbeArm::Cpu => {}
1009                        }
1010                    }
1011                    q4t_matvec(self.quant_bytes(), x, *rows, *cols, out, pool);
1012                    return;
1013                }
1014                if *dtype == TensorDtype::Q4TiledP {
1015                    // GPU route for large q4tp matvecs — the lm_head class.
1016                    // On a q4tp checkpoint the head is the biggest single
1017                    // host matvec left in the decode step, and the batched
1018                    // kernel at b=1 already exists on both backends. Probe
1019                    // keeps the winner, same as q4_block above.
1020                    if *rows * *cols >= 8_388_608 && crate::gpu::enabled_here() {
1021                        let t0 = std::time::Instant::now();
1022                        let cls = crate::gpu::matvec_class(*rows, *cols);
1023                        match crate::gpu::probe_arm(cls) {
1024                            crate::gpu::ProbeArm::Gpu => {
1025                                if crate::gpu::q4tp_matvec(model, *idx, x, *rows, *cols, out) {
1026                                    crate::gpu::probe_record(cls, true, t0.elapsed());
1027                                    return;
1028                                }
1029                            }
1030                            crate::gpu::ProbeArm::CpuTimed => {
1031                                q4tp_matvec(self.quant_bytes(), x, *rows, *cols, out, pool);
1032                                crate::gpu::probe_record(cls, false, t0.elapsed());
1033                                return;
1034                            }
1035                            crate::gpu::ProbeArm::Cpu => {}
1036                        }
1037                    }
1038                    q4tp_matvec(self.quant_bytes(), x, *rows, *cols, out, pool);
1039                    return;
1040                }
1041                if *dtype == TensorDtype::Q2TiledP {
1042                    q2tp_matvec(self.quant_bytes(), x, *rows, *cols, out, pool);
1043                    return;
1044                }
1045                if *dtype == TensorDtype::Q1 {
1046                    // GPU route for large q1 matvecs (out_proj / lm_head
1047                    // class): the CPU q1 kernel is load-port-bound at
1048                    // ~4 GB/s/core, the GPU one is bandwidth-bound — the
1049                    // probe measures both arms and keeps the winner.
1050                    if *rows * *cols >= 8_388_608 && crate::gpu::enabled_here() {
1051                        let t0 = std::time::Instant::now();
1052                        let arm = if crate::gpu::q1_force() {
1053                            crate::gpu::ProbeArm::Gpu
1054                        } else {
1055                            crate::gpu::probe_arm(crate::gpu::OpClass::Matvec)
1056                        };
1057                        match arm {
1058                            crate::gpu::ProbeArm::Gpu => {
1059                                if crate::gpu::q1_matvec(model, *idx, x, *rows, *cols, out) {
1060                                    crate::gpu::probe_record(
1061                                        crate::gpu::OpClass::Matvec,
1062                                        true,
1063                                        t0.elapsed(),
1064                                    );
1065                                    return;
1066                                }
1067                            }
1068                            crate::gpu::ProbeArm::CpuTimed => {
1069                                q1_matvec(self.quant_bytes(), x, *rows, *cols, out, pool);
1070                                crate::gpu::probe_record(
1071                                    crate::gpu::OpClass::Matvec,
1072                                    false,
1073                                    t0.elapsed(),
1074                                );
1075                                return;
1076                            }
1077                            crate::gpu::ProbeArm::Cpu => {}
1078                        }
1079                    }
1080                    q1_matvec(self.quant_bytes(), x, *rows, *cols, out, pool);
1081                    return;
1082                }
1083                if *dtype == TensorDtype::Q1T {
1084                    // GPU route for large q1t matvecs: the ternary BASE dot runs
1085                    // on the GPU (load-port-bound on CPU, like q1), then the
1086                    // sparse overlay is added on the CPU. Probe keeps the winner.
1087                    if *rows * *cols >= 8_388_608 && crate::gpu::enabled_here() {
1088                        let t0 = std::time::Instant::now();
1089                        match crate::gpu::probe_arm(crate::gpu::OpClass::Matvec) {
1090                            crate::gpu::ProbeArm::Gpu => {
1091                                if crate::gpu::q1t_matvec(model, *idx, x, *rows, *cols, out) {
1092                                    q1t_add_overlay(self.quant_bytes(), x, *rows, *cols, out, pool);
1093                                    crate::gpu::probe_record(
1094                                        crate::gpu::OpClass::Matvec,
1095                                        true,
1096                                        t0.elapsed(),
1097                                    );
1098                                    return;
1099                                }
1100                            }
1101                            crate::gpu::ProbeArm::CpuTimed => {
1102                                q1t_matvec(self.quant_bytes(), x, *rows, *cols, out, pool);
1103                                crate::gpu::probe_record(
1104                                    crate::gpu::OpClass::Matvec,
1105                                    false,
1106                                    t0.elapsed(),
1107                                );
1108                                return;
1109                            }
1110                            crate::gpu::ProbeArm::Cpu => {}
1111                        }
1112                    }
1113                    q1t_matvec(self.quant_bytes(), x, *rows, *cols, out, pool);
1114                    return;
1115                }
1116                if matches!(dtype, TensorDtype::Vbit | TensorDtype::VbitRo) {
1117                    vbitmatvec(self.quant_bytes(), vbit_offsets, x, *rows, *cols, out, pool);
1118                    return;
1119                }
1120                let xs = prescale(x, col_field, *dtype);
1121                // D5: large q8 matrices (lm_head-class) — hybrid
1122                // CPU∥GPU: split the rows, both sides compute
1123                // SIMULTANEOUSLY (same math, shared prescale).
1124                // GPU share: CMF_GPU_SPLIT (0..1, default 0.5).
1125                if *rows >= crate::gpu::min_rows()
1126                    && matches!(dtype, TensorDtype::Q8Row | TensorDtype::Q8_2f)
1127                    && gpu_lmhead_enabled()
1128                    && crate::gpu::enabled_here()
1129                {
1130                    // Runtime probe: alternate the hybrid against the
1131                    // pure-CPU matvec, keep whichever is faster HERE.
1132                    let t0 = std::time::Instant::now();
1133                    match crate::gpu::probe_arm(crate::gpu::OpClass::Matvec) {
1134                        crate::gpu::ProbeArm::Gpu => {}
1135                        crate::gpu::ProbeArm::CpuTimed => {
1136                            qmatvec(
1137                                self.quant_bytes(),
1138                                repack,
1139                                row_scale,
1140                                x,
1141                                col_field,
1142                                *dtype,
1143                                *rows,
1144                                *cols,
1145                                out,
1146                                pool,
1147                            );
1148                            crate::gpu::probe_record(
1149                                crate::gpu::OpClass::Matvec,
1150                                false,
1151                                t0.elapsed(),
1152                            );
1153                            return;
1154                        }
1155                        crate::gpu::ProbeArm::Cpu => {
1156                            qmatvec(
1157                                self.quant_bytes(),
1158                                repack,
1159                                row_scale,
1160                                x,
1161                                col_field,
1162                                *dtype,
1163                                *rows,
1164                                *cols,
1165                                out,
1166                                pool,
1167                            );
1168                            return;
1169                        }
1170                    }
1171                    let frac = gpu_split_frac();
1172                    let cpu_rows = ((*rows as f32) * (1.0 - frac)) as usize;
1173                    let (out_cpu, out_gpu) = out.split_at_mut(cpu_rows);
1174                    let bytes = self.quant_bytes();
1175                    let ok = std::thread::scope(|sc| {
1176                        let g = sc.spawn(|| {
1177                            crate::gpu::q8_matvec_range(
1178                                model,
1179                                *idx,
1180                                cpu_rows,
1181                                &row_scale[cpu_rows..],
1182                                &xs,
1183                                *rows - cpu_rows,
1184                                *cols,
1185                                out_gpu,
1186                            )
1187                        });
1188                        if cpu_rows > 0 {
1189                            // Repack prefix covers the full groups of the
1190                            // CPU half (the split starts at row 0).
1191                            let rep_cpu = if repack.is_empty() {
1192                                &[][..]
1193                            } else {
1194                                &repack[..(cpu_rows / 4) * 4 * *cols]
1195                            };
1196                            qmatvec(
1197                                &bytes[..cpu_rows * *cols],
1198                                rep_cpu,
1199                                &row_scale[..cpu_rows],
1200                                x,
1201                                col_field,
1202                                *dtype,
1203                                cpu_rows,
1204                                *cols,
1205                                out_cpu,
1206                                pool,
1207                            );
1208                        }
1209                        g.join().unwrap_or(false)
1210                    });
1211                    if ok {
1212                        crate::gpu::probe_record(crate::gpu::OpClass::Matvec, true, t0.elapsed());
1213                        return;
1214                    }
1215                    // GPU failed — CPU finishes its half (rows rebased —
1216                    // group offsets don't line up, mmap layout only).
1217                    qmatvec(
1218                        &bytes[cpu_rows * *cols..(*rows) * *cols],
1219                        &[],
1220                        &row_scale[cpu_rows..],
1221                        x,
1222                        col_field,
1223                        *dtype,
1224                        *rows - cpu_rows,
1225                        *cols,
1226                        out_gpu,
1227                        pool,
1228                    );
1229                    return;
1230                }
1231                qmatvec(
1232                    self.quant_bytes(),
1233                    repack,
1234                    row_scale,
1235                    x,
1236                    col_field,
1237                    *dtype,
1238                    *rows,
1239                    *cols,
1240                    out,
1241                    pool,
1242                );
1243            }
1244        }
1245    }
1246
1247    /// Fused two-input matvec (MTP verify pair): weights streamed once.
1248    pub fn matvec2(
1249        &self,
1250        x1: &[f32],
1251        x2: &[f32],
1252        o1: &mut [f32],
1253        o2: &mut [f32],
1254        pool: Option<&Pool>,
1255    ) {
1256        match self {
1257            Self::F32 { data, .. } => matvec_rows2(pool, data, x1, x2, o1, o2),
1258            Self::Mapped {
1259                dtype,
1260                rows,
1261                cols,
1262                row_scale,
1263                col_field,
1264                vbit_offsets,
1265                ..
1266            } => {
1267                if *dtype == TensorDtype::Q4Block {
1268                    q4matvec2(self.quant_bytes(), x1, x2, *rows, *cols, o1, o2, pool);
1269                    return;
1270                }
1271                if *dtype == TensorDtype::Q4Tiled {
1272                    q4t_matvec2(self.quant_bytes(), x1, x2, *rows, *cols, o1, o2, pool);
1273                    return;
1274                }
1275                if *dtype == TensorDtype::Q4TiledP {
1276                    q4tp_matvec2(self.quant_bytes(), x1, x2, *rows, *cols, o1, o2, pool);
1277                    return;
1278                }
1279                if *dtype == TensorDtype::Q2TiledP {
1280                    q2tp_matvec2(self.quant_bytes(), x1, x2, *rows, *cols, o1, o2, pool);
1281                    return;
1282                }
1283                if *dtype == TensorDtype::Q1 {
1284                    q1_matvec2(self.quant_bytes(), x1, x2, *rows, *cols, o1, o2, pool);
1285                    return;
1286                }
1287                if *dtype == TensorDtype::Q1T {
1288                    // Fused ternary pair: one row pass, the register
1289                    // unpack shared across both streams on ARM. (Q1T
1290                    // lacks a row_scale array — scales live inline in
1291                    // the tiles — so it must not fall through to the
1292                    // q8 qmatvec2 below.)
1293                    q1t_matvec2(self.quant_bytes(), x1, x2, *rows, *cols, o1, o2, pool);
1294                    return;
1295                }
1296                if matches!(dtype, TensorDtype::Vbit | TensorDtype::VbitRo) {
1297                    vbitmatvec2(
1298                        self.quant_bytes(),
1299                        vbit_offsets,
1300                        x1,
1301                        x2,
1302                        *rows,
1303                        *cols,
1304                        o1,
1305                        o2,
1306                        pool,
1307                    );
1308                    return;
1309                }
1310                qmatvec2(
1311                    self.quant_bytes(),
1312                    row_scale,
1313                    x1,
1314                    x2,
1315                    col_field,
1316                    *dtype,
1317                    *rows,
1318                    *cols,
1319                    o1,
1320                    o2,
1321                    pool,
1322                );
1323            }
1324        }
1325    }
1326}
1327
1328impl QTensor {
1329    /// Batched matvec (prefill-GEMM): xs — row-major [b, cols],
1330    /// out — row-major [b, rows]. Element-wise semantics are IDENTICAL
1331    /// to b matvec calls (same dot kernels in the same order); the win —
1332    /// the weight row streams from DRAM once per batch, not b times.
1333    /// `(model, index)` when this is a memory-mapped q4tp tensor — the
1334    /// identity a device-resident chain needs to hand `tp_matmat` the
1335    /// weight without going through this struct's own dispatch.
1336    pub fn q4tp_mapped(&self) -> Option<(&std::sync::Arc<CmfModel>, usize)> {
1337        match self {
1338            Self::Mapped {
1339                model, idx, dtype, ..
1340            } if *dtype == TensorDtype::Q4TiledP => Some((model, *idx)),
1341            _ => None,
1342        }
1343    }
1344
1345    pub fn matmat(&self, xs_all: &[f32], b: usize, out: &mut [f32], pool: Option<&Pool>) {
1346        let cols = self.cols();
1347        let rows = self.rows();
1348        debug_assert_eq!(xs_all.len(), b * cols);
1349        debug_assert_eq!(out.len(), b * rows);
1350        // GPTQ calibration: fold this layer's inputs into its Hessian. Only
1351        // Mapped tensors carry a directory name; the check is a relaxed
1352        // atomic load, free when not calibrating.
1353        if crate::gptq_capture::capturing() {
1354            if let Self::Mapped { model, idx, .. } = self {
1355                crate::gptq_capture::accumulate(&model.tensors[*idx].name, xs_all, b, cols);
1356            }
1357        }
1358        match self {
1359            Self::F32 { data, .. } => {
1360                let out_addr = SendMut(out.as_mut_ptr());
1361                let run = |start: usize, end: usize| {
1362                    for o in start..end {
1363                        let row = &data[o * cols..(o + 1) * cols];
1364                        for bi in 0..b {
1365                            let x = &xs_all[bi * cols..(bi + 1) * cols];
1366                            let mut acc = 0f32;
1367                            for j in 0..cols {
1368                                acc += row[j] * x[j];
1369                            }
1370                            unsafe { *out_addr.at(bi * rows + o) = acc };
1371                        }
1372                    }
1373                };
1374                dispatch_rows(pool, rows, &run);
1375            }
1376            Self::Mapped {
1377                dtype,
1378                row_scale,
1379                col_field,
1380                vbit_offsets,
1381                ..
1382            } => {
1383                if *dtype == TensorDtype::Q4Block {
1384                    q4matmat(self.quant_bytes(), xs_all, b, rows, cols, out, pool);
1385                    return;
1386                }
1387                if *dtype == TensorDtype::Q4TiledP {
1388                    // GPU batched q4tp GEMM (dequant + f32nt mul_mm on the
1389                    // device); the probe keeps whichever beats the CPU arm.
1390                    // Narrow (prompt-encode) and wide (DiT) batches probe
1391                    // as separate classes — the regimes have opposite
1392                    // winners and one shared verdict locked the wrong arm.
1393                    // Kill switch (gpu::mm_kill): one grossly slow GPU op
1394                    // (a fair-condition op is ≤~100 ms even at 1024px)
1395                    // means the device is contended by another process
1396                    // (e.g. a simulator) — verdicts are per-process, so
1397                    // without the bail the whole render crawls behind
1398                    // someone else's queue.
1399                    if b >= 32
1400                        && b * rows * cols >= 128_000_000
1401                        && cols % 32 == 0
1402                        && !crate::gpu::mm_killed()
1403                        && crate::gpu::enabled_here()
1404                    {
1405                        let class = if b >= 128 {
1406                            crate::gpu::OpClass::MatmatWide
1407                        } else {
1408                            crate::gpu::OpClass::Matmat
1409                        };
1410                        if let Self::Mapped { model, idx, .. } = self {
1411                            // In-process A/B (`CMF_MM_AB=1`). Three
1412                            // wall-clock A/Bs on a shared stand disagreed
1413                            // with each other by 25% on the same change,
1414                            // because the machine drifts between processes
1415                            // and interleaving whole renders does not fix
1416                            // that. Here both arms run back to back on the
1417                            // SAME data inside one call, so whatever the
1418                            // machine is doing, it does to both — and the
1419                            // disagreement between their outputs falls out
1420                            // for free. Doubles the work; a diagnostic,
1421                            // not a mode.
1422                            if crate::mm_ab::on() {
1423                                let mut g = vec![0f32; b * rows];
1424                                let t = std::time::Instant::now();
1425                                let took = crate::gpu::q4tp_matmat(
1426                                    model, *idx, xs_all, b, rows, cols, &mut g,
1427                                );
1428                                let dg = t.elapsed();
1429                                let t = std::time::Instant::now();
1430                                q4tp_matmat(self.quant_bytes(), xs_all, b, rows, cols, out, pool);
1431                                let dc = t.elapsed();
1432                                crate::mm_ab::record(b, rows, cols, took, dg, dc, &g, out);
1433                                return;
1434                            }
1435                            let t0 = std::time::Instant::now();
1436                            // A cold call takes the device arm: its sample
1437                            // is discarded either way, and the upload is
1438                            // what the next step needs.
1439                            let resident = crate::gpu::weight_is_resident(model, *idx);
1440                            match crate::gpu::probe_arm_cold_prefers_gpu(class, resident) {
1441                                crate::gpu::ProbeArm::Gpu => {
1442                                    if crate::gpu::q4tp_matmat(
1443                                        model, *idx, xs_all, b, rows, cols, out,
1444                                    ) {
1445                                        let el = t0.elapsed();
1446                                        // Work-proportional budget: ~8× the
1447                                        // fair-device estimate (+20 ms slack).
1448                                        // An absolute cap missed the worst
1449                                        // case — contended ops sit at
1450                                        // 100–240 ms each and still bury a
1451                                        // render whose fair op is 3–9 ms.
1452                                        // Cold ops (first PSO build, buffer
1453                                        // alloc) are exempt: a one-off
1454                                        // ~50 ms compile is not contention.
1455                                        let flops = 2.0 * b as f64 * rows as f64 * cols as f64;
1456                                        let budget = std::time::Duration::from_secs_f64(
1457                                            flops / 1.5e12 * 8.0 + 0.020,
1458                                        );
1459                                        crate::gpu::mm_budget_check(
1460                                            "q4tp matmat",
1461                                            el,
1462                                            budget,
1463                                            crate::gpu::probe_was_cold() || !resident,
1464                                        );
1465                                        crate::gpu::probe_record(class, true, el);
1466                                        return;
1467                                    }
1468                                }
1469                                crate::gpu::ProbeArm::CpuTimed => {
1470                                    q4tp_matmat(
1471                                        self.quant_bytes(),
1472                                        xs_all,
1473                                        b,
1474                                        rows,
1475                                        cols,
1476                                        out,
1477                                        pool,
1478                                    );
1479                                    crate::gpu::probe_record(class, false, t0.elapsed());
1480                                    return;
1481                                }
1482                                crate::gpu::ProbeArm::Cpu => {}
1483                            }
1484                        }
1485                    }
1486                    q4tp_matmat(self.quant_bytes(), xs_all, b, rows, cols, out, pool);
1487                    return;
1488                }
1489                if *dtype == TensorDtype::Q2TiledP {
1490                    // Same device arm as q4tp, behind the same probe:
1491                    // the planes differ, the dispatch does not. Without
1492                    // this a q2tp file ran its widest projections on the
1493                    // host while the 4-bit one had the card, which is a
1494                    // codec paying for its size twice.
1495                    if b >= 32
1496                        && b * rows * cols >= 128_000_000
1497                        && cols % 32 == 0
1498                        && !crate::gpu::mm_killed()
1499                        && crate::gpu::enabled_here()
1500                    {
1501                        let class = if b >= 128 {
1502                            crate::gpu::OpClass::MatmatWide
1503                        } else {
1504                            crate::gpu::OpClass::Matmat
1505                        };
1506                        if let Self::Mapped { model, idx, .. } = self {
1507                            let t0 = std::time::Instant::now();
1508                            match crate::gpu::probe_arm(class) {
1509                                crate::gpu::ProbeArm::Gpu => {
1510                                    if crate::gpu::q2tp_matmat(
1511                                        model, *idx, xs_all, b, rows, cols, out,
1512                                    ) {
1513                                        crate::gpu::probe_record(class, true, t0.elapsed());
1514                                        return;
1515                                    }
1516                                }
1517                                crate::gpu::ProbeArm::CpuTimed => {
1518                                    q2tp_matmat(
1519                                        self.quant_bytes(),
1520                                        xs_all,
1521                                        b,
1522                                        rows,
1523                                        cols,
1524                                        out,
1525                                        pool,
1526                                    );
1527                                    crate::gpu::probe_record(class, false, t0.elapsed());
1528                                    return;
1529                                }
1530                                crate::gpu::ProbeArm::Cpu => {}
1531                            }
1532                        }
1533                    }
1534                    // Without a host arm a q2tp tensor falls through to
1535                    // the q8 fallback, which reads it at one BYTE per
1536                    // weight — a 2x overrun that killed pool workers
1537                    // mid-prefill while the dispatcher waited forever.
1538                    q2tp_matmat(self.quant_bytes(), xs_all, b, rows, cols, out, pool);
1539                    return;
1540                }
1541                if *dtype == TensorDtype::Q4Tiled {
1542                    // GPU batched q4t GEMM (dequant + f32nt mul_mm on the
1543                    // device); the probe keeps whichever beats the CPU arm.
1544                    // Narrow (prompt-encode) and wide (DiT) batches probe
1545                    // as separate classes — the regimes have opposite
1546                    // winners and one shared verdict locked the wrong arm.
1547                    // Kill switch (gpu::mm_kill): one grossly slow GPU op
1548                    // (a fair-condition op is ≤~100 ms even at 1024px)
1549                    // means the device is contended by another process
1550                    // (e.g. a simulator) — verdicts are per-process, so
1551                    // without the bail the whole render crawls behind
1552                    // someone else's queue.
1553                    if b >= 32
1554                        && b * rows * cols >= 128_000_000
1555                        && cols % 32 == 0
1556                        && !crate::gpu::mm_killed()
1557                        && crate::gpu::enabled_here()
1558                    {
1559                        let class = if b >= 128 {
1560                            crate::gpu::OpClass::MatmatWide
1561                        } else {
1562                            crate::gpu::OpClass::Matmat
1563                        };
1564                        if let Self::Mapped { model, idx, .. } = self {
1565                            let t0 = std::time::Instant::now();
1566                            match crate::gpu::probe_arm(class) {
1567                                crate::gpu::ProbeArm::Gpu => {
1568                                    if crate::gpu::q4t_matmat(
1569                                        model, *idx, xs_all, b, rows, cols, out,
1570                                    ) {
1571                                        let el = t0.elapsed();
1572                                        // Work-proportional budget: ~8× the
1573                                        // fair-device estimate (+20 ms slack).
1574                                        // An absolute cap missed the worst
1575                                        // case — contended ops sit at
1576                                        // 100–240 ms each and still bury a
1577                                        // render whose fair op is 3–9 ms.
1578                                        // Cold ops (first PSO build, buffer
1579                                        // alloc) are exempt: a one-off
1580                                        // ~50 ms compile is not contention.
1581                                        let flops = 2.0 * b as f64 * rows as f64 * cols as f64;
1582                                        let budget = std::time::Duration::from_secs_f64(
1583                                            flops / 1.5e12 * 8.0 + 0.020,
1584                                        );
1585                                        crate::gpu::mm_budget_check(
1586                                            "q4t matmat",
1587                                            el,
1588                                            budget,
1589                                            crate::gpu::probe_was_cold(),
1590                                        );
1591                                        crate::gpu::probe_record(class, true, el);
1592                                        return;
1593                                    }
1594                                }
1595                                crate::gpu::ProbeArm::CpuTimed => {
1596                                    q4t_matmat(
1597                                        self.quant_bytes(),
1598                                        xs_all,
1599                                        b,
1600                                        rows,
1601                                        cols,
1602                                        out,
1603                                        pool,
1604                                    );
1605                                    crate::gpu::probe_record(class, false, t0.elapsed());
1606                                    return;
1607                                }
1608                                crate::gpu::ProbeArm::Cpu => {}
1609                            }
1610                        }
1611                    }
1612                    q4t_matmat(self.quant_bytes(), xs_all, b, rows, cols, out, pool);
1613                    return;
1614                }
1615                if *dtype == TensorDtype::Q1 {
1616                    // GPU batched q1 GEMM for wide prefill (q1_mul_mm on the
1617                    // device); the probe keeps whichever beats the CPU matmat.
1618                    if b >= 32
1619                        && b * rows * cols >= 128_000_000
1620                        && cols % 64 == 0
1621                        && crate::gpu::enabled_here()
1622                    {
1623                        if let Self::Mapped { model, idx, .. } = self {
1624                            let t0 = std::time::Instant::now();
1625                            match crate::gpu::probe_arm(crate::gpu::OpClass::Matmat) {
1626                                crate::gpu::ProbeArm::Gpu => {
1627                                    if crate::gpu::q1_matmat(
1628                                        model, *idx, xs_all, b, rows, cols, out,
1629                                    ) {
1630                                        crate::gpu::probe_record(
1631                                            crate::gpu::OpClass::Matmat,
1632                                            true,
1633                                            t0.elapsed(),
1634                                        );
1635                                        return;
1636                                    }
1637                                }
1638                                crate::gpu::ProbeArm::CpuTimed => {
1639                                    q1_matmat(self.quant_bytes(), xs_all, b, rows, cols, out, pool);
1640                                    crate::gpu::probe_record(
1641                                        crate::gpu::OpClass::Matmat,
1642                                        false,
1643                                        t0.elapsed(),
1644                                    );
1645                                    return;
1646                                }
1647                                crate::gpu::ProbeArm::Cpu => {}
1648                            }
1649                        }
1650                    }
1651                    q1_matmat(self.quant_bytes(), xs_all, b, rows, cols, out, pool);
1652                    return;
1653                }
1654                if *dtype == TensorDtype::Q1T {
1655                    // GPU batched GEMM for wide prefill (base + overlay on the
1656                    // device); probe keeps the winner vs the CPU matmat.
1657                    if b >= 32 && b * rows * cols >= 128_000_000 && crate::gpu::enabled_here() {
1658                        if let Self::Mapped { model, idx, .. } = self {
1659                            let t0 = std::time::Instant::now();
1660                            match crate::gpu::probe_arm(crate::gpu::OpClass::Matmat) {
1661                                crate::gpu::ProbeArm::Gpu => {
1662                                    if crate::gpu::q1t_matmat(
1663                                        model, *idx, xs_all, b, rows, cols, out,
1664                                    ) {
1665                                        crate::gpu::probe_record(
1666                                            crate::gpu::OpClass::Matmat,
1667                                            true,
1668                                            t0.elapsed(),
1669                                        );
1670                                        return;
1671                                    }
1672                                }
1673                                crate::gpu::ProbeArm::CpuTimed => {
1674                                    q1t_matmat(
1675                                        self.quant_bytes(),
1676                                        xs_all,
1677                                        b,
1678                                        rows,
1679                                        cols,
1680                                        out,
1681                                        pool,
1682                                    );
1683                                    crate::gpu::probe_record(
1684                                        crate::gpu::OpClass::Matmat,
1685                                        false,
1686                                        t0.elapsed(),
1687                                    );
1688                                    return;
1689                                }
1690                                crate::gpu::ProbeArm::Cpu => {}
1691                            }
1692                        }
1693                    }
1694                    q1t_matmat(self.quant_bytes(), xs_all, b, rows, cols, out, pool);
1695                    return;
1696                }
1697                if matches!(dtype, TensorDtype::Vbit | TensorDtype::VbitRo) {
1698                    vbitmatmat(
1699                        self.quant_bytes(),
1700                        vbit_offsets,
1701                        xs_all,
1702                        b,
1703                        rows,
1704                        cols,
1705                        out,
1706                        pool,
1707                    );
1708                    return;
1709                }
1710                let pre: Vec<std::borrow::Cow<'_, [f32]>> = (0..b)
1711                    .map(|bi| prescale(&xs_all[bi * cols..(bi + 1) * cols], col_field, *dtype))
1712                    .collect();
1713                // D5: large prefill-batch GEMMs — on the GPU (threshold by
1714                // work volume: submission carries b×rows×cols MACs).
1715                // Runtime probe: the naive GEMM shader + sync readback
1716                // lose to the CPU GEMM on slow driver stacks — alternate
1717                // both arms and keep the winner.
1718                if b >= 8 && b * rows * cols >= 128_000_000 && crate::gpu::enabled_here() {
1719                    if let Self::Mapped { model, idx, .. } = self {
1720                        let t0 = std::time::Instant::now();
1721                        match crate::gpu::probe_arm(crate::gpu::OpClass::Matmat) {
1722                            crate::gpu::ProbeArm::Gpu
1723                                if crate::gpu::probe_deciding(crate::gpu::OpClass::Matmat)
1724                                    && !crate::gpu::q8_resident_or_upload(model, *idx) =>
1725                            {
1726                                // Cold weights during probing: the upload
1727                                // has started, the count runs on the CPU —
1728                                // the GPU arm samples on the next touch.
1729                                let q = self.quant_bytes();
1730                                qmatmat(q, row_scale, &pre, rows, cols, out, pool);
1731                                return;
1732                            }
1733                            crate::gpu::ProbeArm::Gpu => {
1734                                let flat: Vec<f32> =
1735                                    pre.iter().flat_map(|v| v.iter().copied()).collect();
1736                                if crate::gpu::q8_matmat(
1737                                    model, *idx, row_scale, &flat, b, rows, cols, out,
1738                                ) {
1739                                    crate::gpu::probe_record(
1740                                        crate::gpu::OpClass::Matmat,
1741                                        true,
1742                                        t0.elapsed(),
1743                                    );
1744                                    return;
1745                                }
1746                            }
1747                            crate::gpu::ProbeArm::CpuTimed => {
1748                                let q = self.quant_bytes();
1749                                qmatmat(q, row_scale, &pre, rows, cols, out, pool);
1750                                crate::gpu::probe_record(
1751                                    crate::gpu::OpClass::Matmat,
1752                                    false,
1753                                    t0.elapsed(),
1754                                );
1755                                return;
1756                            }
1757                            crate::gpu::ProbeArm::Cpu => {}
1758                        }
1759                    }
1760                }
1761                let q = self.quant_bytes();
1762                qmatmat(q, row_scale, &pre, rows, cols, out, pool);
1763            }
1764        }
1765    }
1766}
1767
1768impl QTensor {
1769    /// Multi-matrix job (roadmap §3 P0): N tensors sharing one input
1770    /// run under a SINGLE pool dispatch — QKV or gate+up cost one
1771    /// barrier instead of N. Per-row math is the exact same kernel as
1772    /// `matvec` (bit-identical outputs); only the dispatch is fused.
1773    /// Falls back to N sequential matvecs when the set is not a uniform
1774    /// q8-family/F32 group or there is no pool.
1775    pub fn matvec_many<const N: usize>(
1776        ts: [&QTensor; N],
1777        x: &[f32],
1778        mut outs: [&mut [f32]; N],
1779        pool: Option<&Pool>,
1780    ) {
1781        let total_rows: usize = ts.iter().map(|t| t.rows()).sum();
1782        let uniform_q8 = ts.iter().all(|t| {
1783            matches!(
1784                t,
1785                Self::Mapped {
1786                    dtype: TensorDtype::Q8Row | TensorDtype::Q8_2f,
1787                    ..
1788                }
1789            )
1790        });
1791        let uniform_f32 = ts.iter().all(|t| matches!(t, Self::F32 { .. }));
1792        let uniform_q4 = ts.iter().all(|t| {
1793            matches!(
1794                t,
1795                Self::Mapped {
1796                    dtype: TensorDtype::Q4Block,
1797                    ..
1798                }
1799            )
1800        });
1801        let uniform_vbit = ts.iter().all(|t| {
1802            matches!(
1803                t,
1804                Self::Mapped {
1805                    dtype: TensorDtype::Vbit | TensorDtype::VbitRo,
1806                    ..
1807                }
1808            )
1809        });
1810        let uniform_q1 = ts.iter().all(|t| {
1811            matches!(
1812                t,
1813                Self::Mapped {
1814                    dtype: TensorDtype::Q1,
1815                    ..
1816                }
1817            )
1818        });
1819        let uniform_q1t = ts.iter().all(|t| {
1820            matches!(
1821                t,
1822                Self::Mapped {
1823                    dtype: TensorDtype::Q1T,
1824                    ..
1825                }
1826            )
1827        });
1828        // q4tp is the skeleton dtype of the big MoE files, and without an arm
1829        // here every projection that shares an input paid its own pool
1830        // barrier: DeepSeek-V4's attention step alone hands this function
1831        // wq_a, wkv and both compressors' pairs off the same hidden state.
1832        let uniform_q4tp = ts.iter().all(|t| {
1833            matches!(
1834                t,
1835                Self::Mapped {
1836                    dtype: TensorDtype::Q4TiledP,
1837                    ..
1838                }
1839            )
1840        }) && ts
1841            .iter()
1842            .all(|t| t.cols() == ts[0].cols() && t.cols() % GROUP_SIZE == 0);
1843        let Some(pool) = pool else {
1844            for (t, o) in ts.iter().zip(outs.iter_mut()) {
1845                t.matvec(x, o, None);
1846            }
1847            return;
1848        };
1849        if total_rows < 256
1850            || !(uniform_q8
1851                || uniform_f32
1852                || uniform_q4
1853                || uniform_vbit
1854                || uniform_q1
1855                || uniform_q1t
1856                || uniform_q4tp)
1857        {
1858            for (t, o) in ts.iter().zip(outs.iter_mut()) {
1859                t.matvec(x, o, Some(pool));
1860            }
1861            return;
1862        }
1863
1864        if uniform_q4tp {
1865            // Every tensor's rows laid end to end in one virtual row space,
1866            // so the whole set is ONE dispatch. The per-row body is the
1867            // `q4tp_matvec` arm verbatim — same activation split, same
1868            // accumulation order — so the outputs are bit-identical to the
1869            // sequential calls this replaces.
1870            let cols = ts[0].cols();
1871            let gpr = cols / GROUP_SIZE;
1872            let views: Vec<Q4tpView> = ts
1873                .iter()
1874                .map(|t| Q4tpView::new(t.quant_bytes(), t.rows(), cols))
1875                .collect();
1876            let rows_of: Vec<usize> = ts.iter().map(|t| t.rows()).collect();
1877            let outs_addr: [SendMut; N] = std::array::from_fn(|i| SendMut(outs[i].as_mut_ptr()));
1878            // flat index -> (which tensor, which of its rows)
1879            let locate = |flat: usize| -> (usize, usize) {
1880                let mut acc = 0;
1881                for (i, &r) in rows_of.iter().enumerate() {
1882                    if flat < acc + r {
1883                        return (i, flat - acc);
1884                    }
1885                    acc += r;
1886                }
1887                (rows_of.len() - 1, 0)
1888            };
1889            let (views, outs_addr) = (&views, &outs_addr);
1890            if a8w8_enabled() {
1891                let act = split_act(x);
1892                let act = &act;
1893                let run = |start: usize, end: usize| {
1894                    let mut sc = vec![0f32; gpr];
1895                    for flat in start..end {
1896                        let (t, r) = locate(flat);
1897                        let v = &views[t];
1898                        v.scales_into(r, gpr, &mut sc);
1899                        let mut acc = dot_q4tp_row_i8(v.nib, r, gpr, &act.xq, &sc) * act.sx;
1900                        for &(j, xv) in &act.outliers {
1901                            let (w, s) = q4tp_outlier(v.nib, r, gpr, j, &sc);
1902                            acc += w * s * xv;
1903                        }
1904                        // SAFETY: one worker owns each (tensor, row) pair.
1905                        unsafe { *outs_addr[t].at(r) = acc };
1906                    }
1907                };
1908                pool.run_rows(total_rows, &run);
1909            } else {
1910                let run = |start: usize, end: usize| {
1911                    let mut sc = vec![0f32; gpr];
1912                    for flat in start..end {
1913                        let (t, r) = locate(flat);
1914                        let v = &views[t];
1915                        v.scales_into(r, gpr, &mut sc);
1916                        // SAFETY: one worker owns each (tensor, row) pair.
1917                        unsafe { *outs_addr[t].at(r) = q4tp_row_exact(v.nib, r, gpr, x, &sc) };
1918                    }
1919                };
1920                pool.run_rows(total_rows, &run);
1921            }
1922            return;
1923        }
1924
1925        if uniform_q1 {
1926            // One shared activation split + group sums (q1 has no col
1927            // field; the same input feeds every tensor).
1928            let outs_addr: [SendMut; N] = std::array::from_fn(|i| SendMut(outs[i].as_mut_ptr()));
1929            if a8w8_enabled() {
1930                let act = split_act(x);
1931                let gsum = q1_group_sums(&act.xq, ts[0].cols() / GROUP_SIZE);
1932                let (act, gsum) = (&act, &gsum);
1933                let closures: [_; N] = std::array::from_fn(|i| {
1934                    let (bytes, gpr, out) =
1935                        (ts[i].quant_bytes(), ts[i].cols() / GROUP_SIZE, outs_addr[i]);
1936                    move |s: usize, e: usize| q1_range_a8w8(bytes, gpr, act, gsum, out, s, e)
1937                });
1938                let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
1939                    std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
1940                pool.run_many(&parts);
1941            } else {
1942                let closures: [_; N] = std::array::from_fn(|i| {
1943                    let (bytes, gpr, out) =
1944                        (ts[i].quant_bytes(), ts[i].cols() / GROUP_SIZE, outs_addr[i]);
1945                    move |s: usize, e: usize| q1_range_f32(bytes, gpr, x, out, s, e)
1946                });
1947                let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
1948                    std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
1949                pool.run_many(&parts);
1950            }
1951            return;
1952        }
1953
1954        if uniform_q1t {
1955            // Q1T batched: one shared activation split + overlay decode,
1956            // all tensors' rows in ONE pool dispatch (saves N−1 dispatches
1957            // and N−1 redundant split_act calls per layer).
1958            let outs_addr: [SendMut; N] = std::array::from_fn(|i| SendMut(outs[i].as_mut_ptr()));
1959            const TILE: usize = cortiq_core::quant::Q1T_TILE;
1960            if a8w8_enabled() {
1961                let act = split_act(x);
1962                let act = &act;
1963                let x_ref = x;
1964                let closures: [_; N] = std::array::from_fn(|i| {
1965                    let bytes = ts[i].quant_bytes();
1966                    let (rows, cols) = (ts[i].rows(), ts[i].cols());
1967                    let gpr = cols / GROUP_SIZE;
1968                    let (rp_off, ent_off, has_ov) = q1t_overlay(bytes, rows * gpr * TILE, rows);
1969                    let out = outs_addr[i];
1970                    move |s: usize, e: usize| {
1971                        q1t_range_a8w8(bytes, gpr, rp_off, ent_off, has_ov, act, x_ref, out, s, e)
1972                    }
1973                });
1974                let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
1975                    std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
1976                pool.run_many(&parts);
1977            } else {
1978                let x_ref = x;
1979                let closures: [_; N] = std::array::from_fn(|i| {
1980                    let bytes = ts[i].quant_bytes();
1981                    let (rows, cols) = (ts[i].rows(), ts[i].cols());
1982                    let gpr = cols / GROUP_SIZE;
1983                    let (rp_off, ent_off, has_ov) = q1t_overlay(bytes, rows * gpr * TILE, rows);
1984                    let out = outs_addr[i];
1985                    move |s: usize, e: usize| {
1986                        q1t_range_f32_batch(bytes, gpr, rp_off, ent_off, has_ov, x_ref, out, s, e)
1987                    }
1988                });
1989                let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
1990                    std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
1991                pool.run_many(&parts);
1992            }
1993            return;
1994        }
1995
1996        if uniform_q4 || uniform_vbit {
1997            let outs_addr: [SendMut; N] = std::array::from_fn(|i| SendMut(outs[i].as_mut_ptr()));
1998            // q4/vbit share one activation split — no per-tensor col field.
1999            if a8w8_enabled() {
2000                let act = split_act(x);
2001                let act = &act;
2002                if uniform_q4 {
2003                    let closures: [_; N] = std::array::from_fn(|i| {
2004                        let (packed, scales) =
2005                            q4_split(ts[i].quant_bytes(), ts[i].rows(), ts[i].cols());
2006                        let (gpr, cols, out) =
2007                            (ts[i].cols() / GROUP_SIZE, ts[i].cols(), outs_addr[i]);
2008                        move |s: usize, e: usize| {
2009                            q4_range_a8w8(packed, scales, gpr, cols, act, out, s, e)
2010                        }
2011                    });
2012                    let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
2013                        std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
2014                    pool.run_many(&parts);
2015                } else {
2016                    let closures: [_; N] = std::array::from_fn(|i| {
2017                        let Self::Mapped { vbit_offsets, .. } = ts[i] else {
2018                            unreachable!()
2019                        };
2020                        let (bytes, rows, cols, out) = (
2021                            ts[i].quant_bytes(),
2022                            ts[i].rows(),
2023                            ts[i].cols(),
2024                            outs_addr[i],
2025                        );
2026                        move |s: usize, e: usize| {
2027                            vbit_range_a8w8(bytes, vbit_offsets, x, act, rows, cols, out, s, e)
2028                        }
2029                    });
2030                    let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
2031                        std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
2032                    pool.run_many(&parts);
2033                }
2034                return;
2035            }
2036            if uniform_q4 {
2037                let closures: [_; N] = std::array::from_fn(|i| {
2038                    let (packed, scales) =
2039                        q4_split(ts[i].quant_bytes(), ts[i].rows(), ts[i].cols());
2040                    let (gpr, out) = (ts[i].cols() / GROUP_SIZE, outs_addr[i]);
2041                    move |s: usize, e: usize| q4_range_f32(packed, scales, gpr, x, out, s, e)
2042                });
2043                let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
2044                    std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
2045                pool.run_many(&parts);
2046            } else {
2047                let closures: [_; N] = std::array::from_fn(|i| {
2048                    let Self::Mapped { vbit_offsets, .. } = ts[i] else {
2049                        unreachable!()
2050                    };
2051                    let (bytes, rows, cols, out) = (
2052                        ts[i].quant_bytes(),
2053                        ts[i].rows(),
2054                        ts[i].cols(),
2055                        outs_addr[i],
2056                    );
2057                    move |s: usize, e: usize| {
2058                        vbit_range_f32(bytes, vbit_offsets, x, rows, cols, out, s, e)
2059                    }
2060                });
2061                let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
2062                    std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
2063                pool.run_many(&parts);
2064            }
2065            return;
2066        }
2067
2068        if uniform_f32 {
2069            let outs_addr: [SendMut; N] = std::array::from_fn(|i| SendMut(outs[i].as_mut_ptr()));
2070            let closures: [_; N] = std::array::from_fn(|i| {
2071                let Self::F32 { data, cols, .. } = ts[i] else {
2072                    unreachable!()
2073                };
2074                let out = outs_addr[i];
2075                move |start: usize, end: usize| {
2076                    for o in start..end {
2077                        let row = &data[o * cols..(o + 1) * cols];
2078                        let mut sum = 0.0f32;
2079                        for j in 0..*cols {
2080                            sum += row[j] * x[j];
2081                        }
2082                        // SAFETY: disjoint (tensor, row) cells per worker.
2083                        unsafe { *out.at(o) = sum };
2084                    }
2085                }
2086            });
2087            let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
2088                std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
2089            pool.run_many(&parts);
2090            return;
2091        }
2092
2093        // Uniform q8-family: per-tensor prescale (q8_2f col fields
2094        // differ per tensor) + the shared range kernels.
2095        struct Ctx<'a> {
2096            bytes: &'a [u8],
2097            #[cfg_attr(not(target_arch = "aarch64"), allow(dead_code))]
2098            rep: &'a [u8],
2099            row_scale: &'a [f32],
2100            cols: usize,
2101            xs: std::borrow::Cow<'a, [f32]>,
2102        }
2103        let ctxs: [Ctx<'_>; N] = std::array::from_fn(|i| {
2104            let Self::Mapped {
2105                dtype,
2106                cols,
2107                row_scale,
2108                col_field,
2109                repack,
2110                ..
2111            } = ts[i]
2112            else {
2113                unreachable!()
2114            };
2115            Ctx {
2116                bytes: ts[i].quant_bytes(),
2117                rep: repack,
2118                row_scale,
2119                cols: *cols,
2120                xs: prescale(x, col_field, *dtype),
2121            }
2122        });
2123        let outs_addr: [SendMut; N] = std::array::from_fn(|i| SendMut(outs[i].as_mut_ptr()));
2124        #[cfg(target_arch = "aarch64")]
2125        if sdot_enabled() {
2126            let acts: [SplitAct; N] = std::array::from_fn(|i| split_act(&ctxs[i].xs));
2127            let closures: [_; N] = std::array::from_fn(|i| {
2128                let (c, act, out) = (&ctxs[i], &acts[i], outs_addr[i]);
2129                move |start: usize, end: usize| {
2130                    q8_range_sdot(c.bytes, c.rep, c.row_scale, act, c.cols, out, start, end)
2131                }
2132            });
2133            let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
2134                std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
2135            pool.run_many(&parts);
2136            return;
2137        }
2138        #[cfg(target_arch = "x86_64")]
2139        if avx2_a8w8_enabled() {
2140            let acts: [SplitAct; N] = std::array::from_fn(|i| split_act(&ctxs[i].xs));
2141            let closures: [_; N] = std::array::from_fn(|i| {
2142                let (c, act, out) = (&ctxs[i], &acts[i], outs_addr[i]);
2143                move |start: usize, end: usize| {
2144                    q8_range_avx2(c.bytes, c.row_scale, act, c.cols, out, start, end)
2145                }
2146            });
2147            let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
2148                std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
2149            pool.run_many(&parts);
2150            return;
2151        }
2152        let closures: [_; N] = std::array::from_fn(|i| {
2153            let (c, out) = (&ctxs[i], outs_addr[i]);
2154            move |start: usize, end: usize| {
2155                q8_range_f32(c.bytes, c.row_scale, &c.xs, c.cols, out, start, end)
2156            }
2157        });
2158        let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
2159            std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
2160        pool.run_many(&parts);
2161    }
2162}
2163
2164impl QTensor {
2165    /// Pair-input multi-matrix job: N tensors × 2 shared inputs under a
2166    /// single pool dispatch — the MTP/pair decode path publishes one job
2167    /// for Q/K/V (and one for gate+up) instead of one per tensor.
2168    /// Per-row math is exactly `matvec2`'s kernels; bit-identical.
2169    #[allow(clippy::needless_range_loop)]
2170    pub fn matvec2_many<const N: usize>(
2171        ts: [&QTensor; N],
2172        x1: &[f32],
2173        x2: &[f32],
2174        mut o1s: [&mut [f32]; N],
2175        mut o2s: [&mut [f32]; N],
2176        pool: Option<&Pool>,
2177    ) {
2178        let total_rows: usize = ts.iter().map(|t| t.rows()).sum();
2179        let uniform_q8 = ts.iter().all(|t| {
2180            matches!(
2181                t,
2182                Self::Mapped {
2183                    dtype: TensorDtype::Q8Row | TensorDtype::Q8_2f,
2184                    ..
2185                }
2186            )
2187        });
2188        let uniform_f32 = ts.iter().all(|t| matches!(t, Self::F32 { .. }));
2189        let uniform_q4 = ts.iter().all(|t| {
2190            matches!(
2191                t,
2192                Self::Mapped {
2193                    dtype: TensorDtype::Q4Block,
2194                    ..
2195                }
2196            )
2197        });
2198        let uniform_vbit = ts.iter().all(|t| {
2199            matches!(
2200                t,
2201                Self::Mapped {
2202                    dtype: TensorDtype::Vbit | TensorDtype::VbitRo,
2203                    ..
2204                }
2205            )
2206        });
2207        let fusable = pool.is_some()
2208            && total_rows >= 256
2209            && (uniform_q8 || uniform_f32 || uniform_q4 || uniform_vbit);
2210        if !fusable {
2211            for i in 0..N {
2212                ts[i].matvec2(x1, x2, o1s[i], o2s[i], pool);
2213            }
2214            return;
2215        }
2216        let pool = pool.unwrap();
2217
2218        if uniform_q4 || uniform_vbit {
2219            let p1: [SendMut; N] = std::array::from_fn(|i| SendMut(o1s[i].as_mut_ptr()));
2220            let p2: [SendMut; N] = std::array::from_fn(|i| SendMut(o2s[i].as_mut_ptr()));
2221            // q4/vbit share activation splits — no per-tensor col field.
2222            if a8w8_enabled() {
2223                let a1 = split_act(x1);
2224                let a2 = split_act(x2);
2225                let (a1, a2) = (&a1, &a2);
2226                if uniform_q4 {
2227                    let closures: [_; N] = std::array::from_fn(|i| {
2228                        let (packed, scales) =
2229                            q4_split(ts[i].quant_bytes(), ts[i].rows(), ts[i].cols());
2230                        let (gpr, cols, o1, o2) =
2231                            (ts[i].cols() / GROUP_SIZE, ts[i].cols(), p1[i], p2[i]);
2232                        move |s: usize, e: usize| {
2233                            q4_range2_a8w8(packed, scales, gpr, cols, a1, a2, o1, o2, s, e)
2234                        }
2235                    });
2236                    let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
2237                        std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
2238                    pool.run_many(&parts);
2239                } else {
2240                    let closures: [_; N] = std::array::from_fn(|i| {
2241                        let Self::Mapped { vbit_offsets, .. } = ts[i] else {
2242                            unreachable!()
2243                        };
2244                        let (bytes, rows, cols, o1, o2) = (
2245                            ts[i].quant_bytes(),
2246                            ts[i].rows(),
2247                            ts[i].cols(),
2248                            p1[i],
2249                            p2[i],
2250                        );
2251                        move |s: usize, e: usize| {
2252                            vbit_range2_a8w8(
2253                                bytes,
2254                                vbit_offsets,
2255                                x1,
2256                                x2,
2257                                a1,
2258                                a2,
2259                                rows,
2260                                cols,
2261                                o1,
2262                                o2,
2263                                s,
2264                                e,
2265                            )
2266                        }
2267                    });
2268                    let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
2269                        std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
2270                    pool.run_many(&parts);
2271                }
2272                return;
2273            }
2274            if uniform_q4 {
2275                let closures: [_; N] = std::array::from_fn(|i| {
2276                    let (packed, scales) =
2277                        q4_split(ts[i].quant_bytes(), ts[i].rows(), ts[i].cols());
2278                    let (gpr, o1, o2) = (ts[i].cols() / GROUP_SIZE, p1[i], p2[i]);
2279                    move |s: usize, e: usize| {
2280                        q4_range2_f32(packed, scales, gpr, x1, x2, o1, o2, s, e)
2281                    }
2282                });
2283                let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
2284                    std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
2285                pool.run_many(&parts);
2286            } else {
2287                let closures: [_; N] = std::array::from_fn(|i| {
2288                    let Self::Mapped { vbit_offsets, .. } = ts[i] else {
2289                        unreachable!()
2290                    };
2291                    let (bytes, rows, cols, o1, o2) = (
2292                        ts[i].quant_bytes(),
2293                        ts[i].rows(),
2294                        ts[i].cols(),
2295                        p1[i],
2296                        p2[i],
2297                    );
2298                    move |s: usize, e: usize| {
2299                        vbit_range2_f32(bytes, vbit_offsets, x1, x2, rows, cols, o1, o2, s, e)
2300                    }
2301                });
2302                let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
2303                    std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
2304                pool.run_many(&parts);
2305            }
2306            return;
2307        }
2308
2309        if uniform_f32 {
2310            let p1: [SendMut; N] = std::array::from_fn(|i| SendMut(o1s[i].as_mut_ptr()));
2311            let p2: [SendMut; N] = std::array::from_fn(|i| SendMut(o2s[i].as_mut_ptr()));
2312            let closures: [_; N] = std::array::from_fn(|i| {
2313                let Self::F32 { data, cols, .. } = ts[i] else {
2314                    unreachable!()
2315                };
2316                let (o1, o2) = (p1[i], p2[i]);
2317                move |start: usize, end: usize| {
2318                    for o in start..end {
2319                        let row = &data[o * cols..(o + 1) * cols];
2320                        let (mut s1, mut s2) = (0.0f32, 0.0f32);
2321                        for j in 0..*cols {
2322                            s1 += row[j] * x1[j];
2323                            s2 += row[j] * x2[j];
2324                        }
2325                        // SAFETY: disjoint (tensor, row) cells per worker.
2326                        unsafe {
2327                            *o1.at(o) = s1;
2328                            *o2.at(o) = s2;
2329                        }
2330                    }
2331                }
2332            });
2333            let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
2334                std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
2335            pool.run_many(&parts);
2336            return;
2337        }
2338
2339        struct Ctx<'a> {
2340            bytes: &'a [u8],
2341            row_scale: &'a [f32],
2342            cols: usize,
2343            xs1: std::borrow::Cow<'a, [f32]>,
2344            xs2: std::borrow::Cow<'a, [f32]>,
2345        }
2346        let ctxs: [Ctx<'_>; N] = std::array::from_fn(|i| {
2347            let Self::Mapped {
2348                dtype,
2349                cols,
2350                row_scale,
2351                col_field,
2352                ..
2353            } = ts[i]
2354            else {
2355                unreachable!()
2356            };
2357            Ctx {
2358                bytes: ts[i].quant_bytes(),
2359                row_scale,
2360                cols: *cols,
2361                xs1: prescale(x1, col_field, *dtype),
2362                xs2: prescale(x2, col_field, *dtype),
2363            }
2364        });
2365        let p1: [SendMut; N] = std::array::from_fn(|i| SendMut(o1s[i].as_mut_ptr()));
2366        let p2: [SendMut; N] = std::array::from_fn(|i| SendMut(o2s[i].as_mut_ptr()));
2367        #[cfg(target_arch = "aarch64")]
2368        if sdot_enabled() {
2369            let acts: [(SplitAct, SplitAct); N] =
2370                std::array::from_fn(|i| (split_act(&ctxs[i].xs1), split_act(&ctxs[i].xs2)));
2371            let closures: [_; N] = std::array::from_fn(|i| {
2372                let (c, a, o1, o2) = (&ctxs[i], &acts[i], p1[i], p2[i]);
2373                move |start: usize, end: usize| {
2374                    q8_range2_sdot(c.bytes, c.row_scale, &a.0, &a.1, c.cols, o1, o2, start, end)
2375                }
2376            });
2377            let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
2378                std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
2379            pool.run_many(&parts);
2380            return;
2381        }
2382        #[cfg(target_arch = "x86_64")]
2383        if avx2_a8w8_enabled() {
2384            let acts: [(SplitAct, SplitAct); N] =
2385                std::array::from_fn(|i| (split_act(&ctxs[i].xs1), split_act(&ctxs[i].xs2)));
2386            let closures: [_; N] = std::array::from_fn(|i| {
2387                let (c, a, o1, o2) = (&ctxs[i], &acts[i], p1[i], p2[i]);
2388                move |start: usize, end: usize| {
2389                    q8_range2_avx2(c.bytes, c.row_scale, &a.0, &a.1, c.cols, o1, o2, start, end)
2390                }
2391            });
2392            let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
2393                std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
2394            pool.run_many(&parts);
2395            return;
2396        }
2397        let closures: [_; N] = std::array::from_fn(|i| {
2398            let (c, o1, o2) = (&ctxs[i], p1[i], p2[i]);
2399            move |start: usize, end: usize| {
2400                q8_range2_f32(
2401                    c.bytes,
2402                    c.row_scale,
2403                    &c.xs1,
2404                    &c.xs2,
2405                    c.cols,
2406                    o1,
2407                    o2,
2408                    start,
2409                    end,
2410                )
2411            }
2412        });
2413        let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
2414            std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
2415        pool.run_many(&parts);
2416    }
2417
2418    /// Fused gate+up matvec with SiLU·mul: for each row r, computes
2419    /// `silu(gate·x) * (up·x)` and writes to `out[r]`. ONE pool dispatch,
2420    /// no intermediate g/u buffers, no separate silu pass. Falls back
2421    /// (returns false) for unsupported dtype combos.
2422    pub fn matvec_silu_mul(
2423        gate: &QTensor,
2424        up: &QTensor,
2425        x: &[f32],
2426        out: &mut [f32],
2427        pool: Option<&Pool>,
2428    ) -> bool {
2429        let inter = gate.rows();
2430        debug_assert_eq!(up.rows(), inter);
2431        debug_assert_eq!(out.len(), inter);
2432        debug_assert_eq!(gate.cols(), up.cols());
2433        if !a8w8_enabled() {
2434            return false;
2435        }
2436        let act = split_act(x);
2437        let act = &act;
2438        let x_ref = x;
2439        let out_addr = SendMut(out.as_mut_ptr());
2440
2441        match (gate, up) {
2442            // Q4Block gate + Q4Block up (most common mobile q4 models)
2443            (
2444                Self::Mapped {
2445                    dtype: TensorDtype::Q4Block,
2446                    ..
2447                },
2448                Self::Mapped {
2449                    dtype: TensorDtype::Q4Block,
2450                    ..
2451                },
2452            ) => {
2453                let (gp, gs) = q4_split(gate.quant_bytes(), gate.rows(), gate.cols());
2454                let (up_p, up_s) = q4_split(up.quant_bytes(), up.rows(), up.cols());
2455                let gpr = gate.cols() / GROUP_SIZE;
2456                let cols = gate.cols();
2457                let run = move |start: usize, end: usize| {
2458                    for r in start..end {
2459                        let mut gv = dot_q4_row_i8(gp, gs, r * gpr, gpr, &act.xq) * act.sx;
2460                        let mut uv = dot_q4_row_i8(up_p, up_s, r * gpr, gpr, &act.xq) * act.sx;
2461                        for &(j, xv) in &act.outliers {
2462                            let flat = r * cols + j;
2463                            let gb = gp[flat / 2];
2464                            let gn = if flat & 1 == 0 { gb & 0x0F } else { gb >> 4 };
2465                            let gsc = f16_to_f32(u16::from_le_bytes([
2466                                gs[(flat / GROUP_SIZE) * 2],
2467                                gs[(flat / GROUP_SIZE) * 2 + 1],
2468                            ]));
2469                            gv += ((gn as i32 - 8) as f32) * gsc * xv;
2470                            let ub = up_p[flat / 2];
2471                            let un = if flat & 1 == 0 { ub & 0x0F } else { ub >> 4 };
2472                            let usc = f16_to_f32(u16::from_le_bytes([
2473                                up_s[(flat / GROUP_SIZE) * 2],
2474                                up_s[(flat / GROUP_SIZE) * 2 + 1],
2475                            ]));
2476                            uv += ((un as i32 - 8) as f32) * usc * xv;
2477                        }
2478                        let silu_g = gv / (1.0 + (-gv).exp());
2479                        // SAFETY: disjoint row ranges per worker.
2480                        unsafe { *out_addr.at(r) = silu_g * uv };
2481                    }
2482                };
2483                dispatch_rows(pool, inter, &run);
2484                true
2485            }
2486            // Q4Tiled gate + Q4Tiled up — one row pass, both tile
2487            // streams sequential, silu·mul fused (same per-row math as
2488            // `q4t_matvec`).
2489            (
2490                Self::Mapped {
2491                    dtype: TensorDtype::Q4Tiled,
2492                    ..
2493                },
2494                Self::Mapped {
2495                    dtype: TensorDtype::Q4Tiled,
2496                    ..
2497                },
2498            ) => {
2499                let g_bytes = gate.quant_bytes();
2500                let u_bytes = up.quant_bytes();
2501                let gpr = gate.cols() / GROUP_SIZE;
2502                let run = move |start: usize, end: usize| {
2503                    for r in start..end {
2504                        let mut gv = dot_q4t_row_i8(g_bytes, r, gpr, &act.xq) * act.sx;
2505                        let mut uv = dot_q4t_row_i8(u_bytes, r, gpr, &act.xq) * act.sx;
2506                        for &(j, xv) in &act.outliers {
2507                            let (w, s) = q4t_outlier(g_bytes, r, gpr, j);
2508                            gv += w * s * xv;
2509                            let (w, s) = q4t_outlier(u_bytes, r, gpr, j);
2510                            uv += w * s * xv;
2511                        }
2512                        let silu_g = gv / (1.0 + (-gv).exp());
2513                        // SAFETY: disjoint row ranges per worker.
2514                        unsafe { *out_addr.at(r) = silu_g * uv };
2515                    }
2516                };
2517                dispatch_rows(pool, inter, &run);
2518                true
2519            }
2520            // Q4TiledP gate + Q4TiledP up — the same fused row pass, with
2521            // each row's two ladders built once and spent on both streams.
2522            (
2523                Self::Mapped {
2524                    dtype: TensorDtype::Q4TiledP,
2525                    ..
2526                },
2527                Self::Mapped {
2528                    dtype: TensorDtype::Q4TiledP,
2529                    ..
2530                },
2531            ) => {
2532                let cols = gate.cols();
2533                let gpr = cols / GROUP_SIZE;
2534                let gv_view = Q4tpView::new(gate.quant_bytes(), inter, cols);
2535                let uv_view = Q4tpView::new(up.quant_bytes(), inter, cols);
2536                let run = |start: usize, end: usize| {
2537                    let (mut gsc, mut usc) = (vec![0f32; gpr], vec![0f32; gpr]);
2538                    for r in start..end {
2539                        gv_view.scales_into(r, gpr, &mut gsc);
2540                        uv_view.scales_into(r, gpr, &mut usc);
2541                        let mut gv = dot_q4tp_row_i8(gv_view.nib, r, gpr, &act.xq, &gsc) * act.sx;
2542                        let mut uv = dot_q4tp_row_i8(uv_view.nib, r, gpr, &act.xq, &usc) * act.sx;
2543                        for &(j, xv) in &act.outliers {
2544                            let (w, s) = q4tp_outlier(gv_view.nib, r, gpr, j, &gsc);
2545                            gv += w * s * xv;
2546                            let (w, s) = q4tp_outlier(uv_view.nib, r, gpr, j, &usc);
2547                            uv += w * s * xv;
2548                        }
2549                        let silu_g = gv / (1.0 + (-gv).exp());
2550                        // SAFETY: disjoint row ranges per worker.
2551                        unsafe { *out_addr.at(r) = silu_g * uv };
2552                    }
2553                };
2554                dispatch_rows(pool, inter, &run);
2555                true
2556            }
2557            // Q1 gate + Q1 up — one row pass over both sign streams,
2558            // silu·mul fused (the per-row math of `q1_range_a8w8`); the
2559            // activation group sums are shared by both streams. Without
2560            // this arm a q1 dense FFN paid two dispatches + a combine
2561            // loop — the exact barrier this function exists to remove.
2562            (
2563                Self::Mapped {
2564                    dtype: TensorDtype::Q1,
2565                    ..
2566                },
2567                Self::Mapped {
2568                    dtype: TensorDtype::Q1,
2569                    ..
2570                },
2571            ) => {
2572                let g_bytes = gate.quant_bytes();
2573                let u_bytes = up.quant_bytes();
2574                let gpr = gate.cols() / GROUP_SIZE;
2575                let gsum = q1_group_sums(&act.xq, gpr);
2576                let gsum = &gsum;
2577                let run = move |start: usize, end: usize| {
2578                    for r in start..end {
2579                        let mut gv = dot_q1_row_i8(g_bytes, r, gpr, &act.xq, gsum) * act.sx;
2580                        let mut uv = dot_q1_row_i8(u_bytes, r, gpr, &act.xq, gsum) * act.sx;
2581                        for &(j, xv) in &act.outliers {
2582                            let (w, s) = q1_outlier(g_bytes, r, gpr, j);
2583                            gv += w * s * xv;
2584                            let (w, s) = q1_outlier(u_bytes, r, gpr, j);
2585                            uv += w * s * xv;
2586                        }
2587                        let silu_g = gv / (1.0 + (-gv).exp());
2588                        // SAFETY: disjoint row ranges per worker.
2589                        unsafe { *out_addr.at(r) = silu_g * uv };
2590                    }
2591                };
2592                dispatch_rows(pool, inter, &run);
2593                true
2594            }
2595            // Q2TiledP gate + Q2TiledP up — the 2-bit expert pair (MoE
2596            // FFNs of the W2 class): one row pass, both ladders built
2597            // once, integer code dots with shared group sums.
2598            (
2599                Self::Mapped {
2600                    dtype: TensorDtype::Q2TiledP,
2601                    ..
2602                },
2603                Self::Mapped {
2604                    dtype: TensorDtype::Q2TiledP,
2605                    ..
2606                },
2607            ) => {
2608                let cols = gate.cols();
2609                let gpr = cols / GROUP_SIZE;
2610                let gv_view = Q4tpView::new_q2(gate.quant_bytes(), inter, cols);
2611                let uv_view = Q4tpView::new_q2(up.quant_bytes(), inter, cols);
2612                let gsum = q1_group_sums(&act.xq, gpr);
2613                let gsum = &gsum;
2614                let run = move |start: usize, end: usize| {
2615                    let (mut gsc, mut usc) = (vec![0f32; gpr], vec![0f32; gpr]);
2616                    for r in start..end {
2617                        gv_view.scales_into(r, gpr, &mut gsc);
2618                        uv_view.scales_into(r, gpr, &mut usc);
2619                        let mut gv =
2620                            dot_q2tp_row_i8(gv_view.nib, r, gpr, &act.xq, gsum, &gsc) * act.sx;
2621                        let mut uv =
2622                            dot_q2tp_row_i8(uv_view.nib, r, gpr, &act.xq, gsum, &usc) * act.sx;
2623                        for &(j, xv) in &act.outliers {
2624                            let (w, s) = q2tp_outlier(gv_view.nib, r, gpr, j, &gsc);
2625                            gv += w * s * xv;
2626                            let (w, s) = q2tp_outlier(uv_view.nib, r, gpr, j, &usc);
2627                            uv += w * s * xv;
2628                        }
2629                        let silu_g = gv / (1.0 + (-gv).exp());
2630                        // SAFETY: disjoint row ranges per worker.
2631                        unsafe { *out_addr.at(r) = silu_g * uv };
2632                    }
2633                };
2634                dispatch_rows(pool, inter, &run);
2635                true
2636            }
2637            // Q8Row gate + Q8Row up — one row pass over both i8 streams.
2638            // Q8_2f stays out on purpose: its column field prescales the
2639            // activations PER TENSOR, which breaks this fn's shared
2640            // split_act contract — it keeps the two-dispatch path.
2641            (
2642                Self::Mapped {
2643                    dtype: TensorDtype::Q8Row,
2644                    row_scale: g_rs,
2645                    ..
2646                },
2647                Self::Mapped {
2648                    dtype: TensorDtype::Q8Row,
2649                    row_scale: u_rs,
2650                    ..
2651                },
2652            ) => {
2653                let g_bytes = gate.quant_bytes();
2654                let u_bytes = up.quant_bytes();
2655                let cols = gate.cols();
2656                let run = move |start: usize, end: usize| {
2657                    for r in start..end {
2658                        let gv = q8_row_dot(&g_bytes[r * cols..(r + 1) * cols], act) * g_rs[r];
2659                        let uv = q8_row_dot(&u_bytes[r * cols..(r + 1) * cols], act) * u_rs[r];
2660                        let silu_g = gv / (1.0 + (-gv).exp());
2661                        // SAFETY: disjoint row ranges per worker.
2662                        unsafe { *out_addr.at(r) = silu_g * uv };
2663                    }
2664                };
2665                dispatch_rows(pool, inter, &run);
2666                true
2667            }
2668            // Q1T gate + Q1T up
2669            (
2670                Self::Mapped {
2671                    dtype: TensorDtype::Q1T,
2672                    ..
2673                },
2674                Self::Mapped {
2675                    dtype: TensorDtype::Q1T,
2676                    ..
2677                },
2678            ) => {
2679                const TILE: usize = cortiq_core::quant::Q1T_TILE;
2680                let g_bytes = gate.quant_bytes();
2681                let u_bytes = up.quant_bytes();
2682                let gpr = gate.cols() / GROUP_SIZE;
2683                let (g_rp, g_ent, g_ov) = q1t_overlay(g_bytes, inter * gpr * TILE, inter);
2684                let (u_rp, u_ent, u_ov) = q1t_overlay(u_bytes, inter * gpr * TILE, inter);
2685                let run = move |start: usize, end: usize| {
2686                    for r in start..end {
2687                        let mut gv = q1t_dot_row_i8(g_bytes, r, gpr, &act.xq) * act.sx;
2688                        let mut uv = q1t_dot_row_i8(u_bytes, r, gpr, &act.xq) * act.sx;
2689                        for &(j, xv) in &act.outliers {
2690                            gv += q1t_base_weight(g_bytes, r, gpr, j) * xv;
2691                            uv += q1t_base_weight(u_bytes, r, gpr, j) * xv;
2692                        }
2693                        gv += q1t_row_outlier_correction(g_bytes, r, g_rp, g_ent, g_ov, x_ref);
2694                        uv += q1t_row_outlier_correction(u_bytes, r, u_rp, u_ent, u_ov, x_ref);
2695                        let silu_g = gv / (1.0 + (-gv).exp());
2696                        // SAFETY: disjoint row ranges per worker.
2697                        unsafe { *out_addr.at(r) = silu_g * uv };
2698                    }
2699                };
2700                dispatch_rows(pool, inter, &run);
2701                true
2702            }
2703            _ => false,
2704        }
2705    }
2706
2707    /// Every routed expert's fused gate/up/SiLU under ONE pool dispatch.
2708    ///
2709    /// The per-expert path pays a pool barrier per expert per stage: at 9
2710    /// experts over 40 layers that is ~720 barriers a token, and a decode
2711    /// profile of Qwen3.6-35B-A3B showed the pool parked in
2712    /// `psynch_cvwait` about twice as long as it spent computing. Laying
2713    /// every expert's rows end-to-end in one virtual row space collapses
2714    /// the stage to a single dispatch. The per-row body is the
2715    /// single-expert q4tp arm verbatim, so outputs are bit-identical.
2716    ///
2717    /// `false` = something is outside the fused q4tp kernel (dtype, shape,
2718    /// or the `CMF_SDOT=0` exact contract); the caller walks the ordinary
2719    /// per-expert path.
2720    pub fn moe_gate_up_many(
2721        pairs: &[(&QTensor, &QTensor)],
2722        x: &[f32],
2723        outs: &mut [Vec<f32>],
2724        pool: Option<&Pool>,
2725    ) -> bool {
2726        if pairs.is_empty() || pairs.len() != outs.len() || !a8w8_enabled() {
2727            return false;
2728        }
2729        let inter = pairs[0].0.rows();
2730        let cols = pairs[0].0.cols();
2731        if cols % GROUP_SIZE != 0 {
2732            return false;
2733        }
2734        let gpr = cols / GROUP_SIZE;
2735        // Uniform layout across every routed pair: q4tp, or the 2-bit
2736        // profile's q2tp gate/up (the W2 class). Mixed sets refuse.
2737        let q2 = matches!(
2738            pairs[0].0,
2739            Self::Mapped {
2740                dtype: TensorDtype::Q2TiledP,
2741                ..
2742            }
2743        );
2744        let want = if q2 {
2745            TensorDtype::Q2TiledP
2746        } else {
2747            TensorDtype::Q4TiledP
2748        };
2749        let mut views = Vec::with_capacity(pairs.len() * 2);
2750        for ((g, u), o) in pairs.iter().zip(outs.iter()) {
2751            let both = matches!(g, Self::Mapped { dtype, .. } if *dtype == want)
2752                && matches!(u, Self::Mapped { dtype, .. } if *dtype == want);
2753            if !both
2754                || g.rows() != inter
2755                || u.rows() != inter
2756                || g.cols() != cols
2757                || u.cols() != cols
2758                || o.len() != inter
2759            {
2760                return false;
2761            }
2762            let mk = if q2 { Q4tpView::new_q2 } else { Q4tpView::new };
2763            views.push(mk(g.quant_bytes(), inter, cols));
2764            views.push(mk(u.quant_bytes(), inter, cols));
2765        }
2766        let act = split_act(x);
2767        let gsum = if q2 {
2768            q1_group_sums(&act.xq, gpr)
2769        } else {
2770            Vec::new()
2771        };
2772        let (act, gsum) = (&act, &gsum);
2773        let ptrs: Vec<SendMut> = outs.iter_mut().map(|o| SendMut(o.as_mut_ptr())).collect();
2774        let (views, ptrs) = (&views, &ptrs);
2775        let run = |start: usize, end: usize| {
2776            let (mut gsc, mut usc) = (vec![0f32; gpr], vec![0f32; gpr]);
2777            for flat in start..end {
2778                let (e, r) = (flat / inter, flat % inter);
2779                let gv_view = &views[e * 2];
2780                let uv_view = &views[e * 2 + 1];
2781                gv_view.scales_into(r, gpr, &mut gsc);
2782                uv_view.scales_into(r, gpr, &mut usc);
2783                let (mut gv, mut uv) = if q2 {
2784                    (
2785                        dot_q2tp_row_i8(gv_view.nib, r, gpr, &act.xq, gsum, &gsc) * act.sx,
2786                        dot_q2tp_row_i8(uv_view.nib, r, gpr, &act.xq, gsum, &usc) * act.sx,
2787                    )
2788                } else {
2789                    (
2790                        dot_q4tp_row_i8(gv_view.nib, r, gpr, &act.xq, &gsc) * act.sx,
2791                        dot_q4tp_row_i8(uv_view.nib, r, gpr, &act.xq, &usc) * act.sx,
2792                    )
2793                };
2794                for &(j, xv) in &act.outliers {
2795                    let (og, ou) = if q2 {
2796                        (
2797                            q2tp_outlier(gv_view.nib, r, gpr, j, &gsc),
2798                            q2tp_outlier(uv_view.nib, r, gpr, j, &usc),
2799                        )
2800                    } else {
2801                        (
2802                            q4tp_outlier(gv_view.nib, r, gpr, j, &gsc),
2803                            q4tp_outlier(uv_view.nib, r, gpr, j, &usc),
2804                        )
2805                    };
2806                    gv += og.0 * og.1 * xv;
2807                    uv += ou.0 * ou.1 * xv;
2808                }
2809                let silu_g = gv / (1.0 + (-gv).exp());
2810                // SAFETY: one worker owns each (expert, row) pair.
2811                unsafe { *ptrs[e].at(r) = silu_g * uv };
2812            }
2813        };
2814        dispatch_rows(pool, pairs.len() * inter, &run);
2815        true
2816    }
2817
2818    /// Every routed expert's down projection, weighted and summed into
2819    /// `out`, under ONE pool dispatch.
2820    ///
2821    /// Partitioned by OUTPUT row rather than by expert: each row is owned
2822    /// by a single worker, so the experts are summed in the caller's order
2823    /// — the same sequence of f32 adds the serial `out[i] += w·eo[i]` loop
2824    /// performs, hence bit-identical. Partitioning by expert instead would
2825    /// race on the shared accumulator.
2826    pub fn moe_down_many(
2827        downs: &[&QTensor],
2828        gs: &[Vec<f32>],
2829        weights: &[f32],
2830        out: &mut [f32],
2831        pool: Option<&Pool>,
2832    ) -> bool {
2833        if downs.is_empty()
2834            || downs.len() != gs.len()
2835            || downs.len() != weights.len()
2836            || !a8w8_enabled()
2837        {
2838            return false;
2839        }
2840        let rows = out.len();
2841        let cols = downs[0].cols();
2842        if cols % GROUP_SIZE != 0 {
2843            return false;
2844        }
2845        let gpr = cols / GROUP_SIZE;
2846        let mut views = Vec::with_capacity(downs.len());
2847        for (d, g) in downs.iter().zip(gs.iter()) {
2848            if !matches!(
2849                d,
2850                Self::Mapped {
2851                    dtype: TensorDtype::Q4TiledP,
2852                    ..
2853                }
2854            ) || d.rows() != rows
2855                || d.cols() != cols
2856                || g.len() != cols
2857            {
2858                return false;
2859            }
2860            views.push(Q4tpView::new(d.quant_bytes(), rows, cols));
2861        }
2862        // One int8 split per expert — the activation vectors differ.
2863        let acts: Vec<SplitAct> = gs.iter().map(|g| split_act(g)).collect();
2864        // Partitioned by OUTPUT row, with the experts folded inside: each
2865        // row is owned by one worker, so they are summed in the caller's
2866        // order — the same f32 sequence the serial `out[i] += w·eo[i]`
2867        // loop produces. Partitioning by expert instead would either race
2868        // on the accumulator or need a scratch plane and a second pass;
2869        // measured, that variant was a wash, so this keeps the simpler
2870        // shape.
2871        let out_addr = SendMut(out.as_mut_ptr());
2872        let (views, acts, weights) = (&views, &acts, &weights);
2873        let run = |start: usize, end: usize| {
2874            let mut sc = vec![0f32; gpr];
2875            for r in start..end {
2876                let mut acc = 0f32;
2877                for (e, v) in views.iter().enumerate() {
2878                    v.scales_into(r, gpr, &mut sc);
2879                    let a = &acts[e];
2880                    let mut d = dot_q4tp_row_i8(v.nib, r, gpr, &a.xq, &sc) * a.sx;
2881                    for &(j, xv) in &a.outliers {
2882                        let (w, s) = q4tp_outlier(v.nib, r, gpr, j, &sc);
2883                        d += w * s * xv;
2884                    }
2885                    acc += weights[e] * d;
2886                }
2887                // SAFETY: disjoint row ranges per worker.
2888                unsafe { *out_addr.at(r) = acc };
2889            }
2890        };
2891        dispatch_rows(pool, rows, &run);
2892        true
2893    }
2894}
2895
2896/// Batched q8 kernel: same math as qmatvec, the row makes a single
2897/// pass from memory for the whole batch.
2898/// Accelerate CBLAS — the Apple AMX matrix units, the same engine
2899/// llama.cpp's `-ngl 0` prefill rides via ggml-blas.
2900#[cfg(target_os = "macos")]
2901mod accel_blas {
2902    #[link(name = "Accelerate", kind = "framework")]
2903    unsafe extern "C" {
2904        pub fn cblas_sgemm(
2905            order: i32,
2906            trans_a: i32,
2907            trans_b: i32,
2908            m: i32,
2909            n: i32,
2910            k: i32,
2911            alpha: f32,
2912            a: *const f32,
2913            lda: i32,
2914            b: *const f32,
2915            ldb: i32,
2916            beta: f32,
2917            c: *mut f32,
2918            ldc: i32,
2919        );
2920    }
2921}
2922
2923#[cfg(target_os = "macos")]
2924pub(crate) fn accel_gemm_enabled() -> bool {
2925    static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
2926    *ON.get_or_init(|| std::env::var("CMF_ACCEL").map(|v| v != "0").unwrap_or(true))
2927}
2928
2929/// Off macOS the "accel" GEMM is the portable NEON micro-kernel below —
2930/// same entry point, so the batched-attention path opens on mobile.
2931#[cfg(all(target_arch = "aarch64", not(target_os = "macos")))]
2932pub(crate) fn accel_gemm_enabled() -> bool {
2933    true
2934}
2935
2936/// Portable NEON f32 GEMM (row-major, optional Bᵀ): a 4×8 fmla
2937/// micro-kernel with A broadcast against B panels — the mobile stand-in
2938/// for Accelerate in the batched causal attention (QKᵀ and P·V). Not a
2939/// BLAS: shapes here are the attention panels (m ≤ heads·chunk,
2940/// k = head_dim or context), and the goal is removing the per-position
2941/// quadratic wall, not peak GEMM.
2942#[cfg(target_arch = "aarch64")]
2943#[allow(clippy::too_many_arguments)]
2944pub(crate) fn neon_gemm_rm(
2945    m: usize,
2946    n: usize,
2947    k: usize,
2948    alpha: f32,
2949    a: &[f32],
2950    lda: usize,
2951    b_mat: &[f32],
2952    ldb: usize,
2953    b_rows_are_n: bool,
2954    c: &mut [f32],
2955    ldc: usize,
2956) {
2957    debug_assert!(a.len() >= (m - 1) * lda + k);
2958    debug_assert!(c.len() >= (m - 1) * ldc + n);
2959    // SAFETY: bounds asserted above; NEON is baseline on aarch64.
2960    unsafe {
2961        use core::arch::aarch64::*;
2962        let mut i = 0usize;
2963        while i < m {
2964            let mi = (m - i).min(4);
2965            let mut j = 0usize;
2966            while j < n {
2967                let nj = (n - j).min(8);
2968                if mi == 4 && nj == 8 {
2969                    let (mut c0a, mut c0b) = (vdupq_n_f32(0.0), vdupq_n_f32(0.0));
2970                    let (mut c1a, mut c1b) = (vdupq_n_f32(0.0), vdupq_n_f32(0.0));
2971                    let (mut c2a, mut c2b) = (vdupq_n_f32(0.0), vdupq_n_f32(0.0));
2972                    let (mut c3a, mut c3b) = (vdupq_n_f32(0.0), vdupq_n_f32(0.0));
2973                    for p in 0..k {
2974                        let (b0, b1) = if b_rows_are_n {
2975                            // B is [n, k]: column p of Bᵀ = element p of
2976                            // eight consecutive B rows — gathered.
2977                            let base = b_mat.as_ptr().add(j * ldb + p);
2978                            let g = |o: usize| *base.add(o * ldb);
2979                            ([g(0), g(1), g(2), g(3)], [g(4), g(5), g(6), g(7)])
2980                        } else {
2981                            let base = b_mat.as_ptr().add(p * ldb + j);
2982                            (
2983                                [*base, *base.add(1), *base.add(2), *base.add(3)],
2984                                [*base.add(4), *base.add(5), *base.add(6), *base.add(7)],
2985                            )
2986                        };
2987                        let bv0 = vld1q_f32(b0.as_ptr());
2988                        let bv1 = vld1q_f32(b1.as_ptr());
2989                        let a0 = vdupq_n_f32(*a.as_ptr().add(i * lda + p));
2990                        let a1 = vdupq_n_f32(*a.as_ptr().add((i + 1) * lda + p));
2991                        let a2 = vdupq_n_f32(*a.as_ptr().add((i + 2) * lda + p));
2992                        let a3 = vdupq_n_f32(*a.as_ptr().add((i + 3) * lda + p));
2993                        c0a = vfmaq_f32(c0a, a0, bv0);
2994                        c0b = vfmaq_f32(c0b, a0, bv1);
2995                        c1a = vfmaq_f32(c1a, a1, bv0);
2996                        c1b = vfmaq_f32(c1b, a1, bv1);
2997                        c2a = vfmaq_f32(c2a, a2, bv0);
2998                        c2b = vfmaq_f32(c2b, a2, bv1);
2999                        c3a = vfmaq_f32(c3a, a3, bv0);
3000                        c3b = vfmaq_f32(c3b, a3, bv1);
3001                    }
3002                    let al = vdupq_n_f32(alpha);
3003                    for (r, (ca, cb)) in [(c0a, c0b), (c1a, c1b), (c2a, c2b), (c3a, c3b)]
3004                        .iter()
3005                        .enumerate()
3006                    {
3007                        let dst = c.as_mut_ptr().add((i + r) * ldc + j);
3008                        vst1q_f32(dst, vmulq_f32(*ca, al));
3009                        vst1q_f32(dst.add(4), vmulq_f32(*cb, al));
3010                    }
3011                } else {
3012                    for r in 0..mi {
3013                        for q in 0..nj {
3014                            let mut acc = 0f32;
3015                            for p in 0..k {
3016                                let bv = if b_rows_are_n {
3017                                    b_mat[(j + q) * ldb + p]
3018                                } else {
3019                                    b_mat[p * ldb + j + q]
3020                                };
3021                                acc += a[(i + r) * lda + p] * bv;
3022                            }
3023                            c[(i + r) * ldc + j + q] = acc * alpha;
3024                        }
3025                    }
3026                }
3027                j += nj;
3028            }
3029            i += mi;
3030        }
3031    }
3032}
3033
3034/// Off-macOS aarch64: the batched attention rides the NEON micro-GEMM.
3035#[cfg(all(target_arch = "aarch64", not(target_os = "macos")))]
3036#[allow(clippy::too_many_arguments)]
3037pub(crate) fn sgemm_rm(
3038    m: usize,
3039    n: usize,
3040    k: usize,
3041    alpha: f32,
3042    a: &[f32],
3043    lda: usize,
3044    b_mat: &[f32],
3045    ldb: usize,
3046    b_rows_are_n: bool,
3047    c: &mut [f32],
3048    ldc: usize,
3049) {
3050    neon_gemm_rm(m, n, k, alpha, a, lda, b_mat, ldb, b_rows_are_n, c, ldc);
3051}
3052
3053/// Row-major f32 GEMM, exposed for offline tools (the AWNP pass builds a
3054/// per-layer projection and applies it to every expert; a naive triple loop
3055/// would turn a two-minute job into half an hour).
3056#[allow(clippy::too_many_arguments)]
3057pub fn sgemm_public(
3058    m: usize,
3059    n: usize,
3060    k: usize,
3061    alpha: f32,
3062    a: &[f32],
3063    lda: usize,
3064    b_mat: &[f32],
3065    ldb: usize,
3066    b_rows_are_n: bool,
3067    c: &mut [f32],
3068    ldc: usize,
3069) {
3070    #[cfg(any(target_os = "macos", target_arch = "aarch64"))]
3071    {
3072        sgemm_rm(m, n, k, alpha, a, lda, b_mat, ldb, b_rows_are_n, c, ldc);
3073    }
3074    // x86 without Accelerate has no sgemm_rm: the specialized paths there are
3075    // quantized kernels, not an f32 GEMM. Only the offline AWNP pass reaches
3076    // this, so correctness matters and throughput does not — a triple loop is
3077    // the honest fallback rather than a reason to make the tool macOS-only.
3078    #[cfg(not(any(target_os = "macos", target_arch = "aarch64")))]
3079    {
3080        for i in 0..m {
3081            for j in 0..n {
3082                let mut acc = 0f32;
3083                for p in 0..k {
3084                    let bv = if b_rows_are_n {
3085                        b_mat[j * ldb + p]
3086                    } else {
3087                        b_mat[p * ldb + j]
3088                    };
3089                    acc += a[i * lda + p] * bv;
3090                }
3091                c[i * ldc + j] = alpha * acc;
3092            }
3093        }
3094    }
3095}
3096
3097/// Row-major f32 GEMM on Accelerate: C[m,n] = alpha·A[m,k] × B(ᵀ).
3098/// `b_rows_are_n` = true multiplies by Bᵀ where B is stored [n, k].
3099#[cfg(target_os = "macos")]
3100#[allow(clippy::too_many_arguments)]
3101pub(crate) fn sgemm_rm(
3102    m: usize,
3103    n: usize,
3104    k: usize,
3105    alpha: f32,
3106    a: &[f32],
3107    lda: usize,
3108    b_mat: &[f32],
3109    ldb: usize,
3110    b_rows_are_n: bool,
3111    c: &mut [f32],
3112    ldc: usize,
3113) {
3114    debug_assert!(a.len() >= (m - 1) * lda + k);
3115    debug_assert!(c.len() >= (m - 1) * ldc + n);
3116    // Test hook: route the attention GEMMs through the portable NEON
3117    // micro-kernel ON APPLE SILICON — how the mobile batched attend is
3118    // measured without a phone in the loop. (Intel macOS has no NEON —
3119    // the hook is a no-op there, Accelerate continues below.)
3120    #[cfg(target_arch = "aarch64")]
3121    if std::env::var("CMF_FORCE_NEON_GEMM")
3122        .map(|v| v == "1")
3123        .unwrap_or(false)
3124    {
3125        return neon_gemm_rm(m, n, k, alpha, a, lda, b_mat, ldb, b_rows_are_n, c, ldc);
3126    }
3127    unsafe {
3128        accel_blas::cblas_sgemm(
3129            101, // RowMajor
3130            111, // NoTrans A
3131            if b_rows_are_n { 112 } else { 111 },
3132            m as i32,
3133            n as i32,
3134            k as i32,
3135            alpha,
3136            a.as_ptr(),
3137            lda as i32,
3138            b_mat.as_ptr(),
3139            ldb as i32,
3140            0.0,
3141            c.as_mut_ptr(),
3142            ldc as i32,
3143        );
3144    }
3145}
3146
3147/// Prefill GEMM through Accelerate (macOS): dequantize q8 rows into
3148/// f32 tiles (scale folded in, pool-parallel) and multiply each tile
3149/// on the AMX with one row-major sgemm. Tiles live in cache, weights
3150/// stream once. Numerics are f32-GEMM (not the int8 dot): prefill
3151/// logits shift within f32 rounding — tolerance-class, like every
3152/// reduction-order change; decode (M=1) never takes this path.
3153#[cfg(target_os = "macos")]
3154fn qmatmat_accel(
3155    q: &[u8],
3156    row_scale: &[f32],
3157    pre: &[std::borrow::Cow<'_, [f32]>],
3158    rows: usize,
3159    cols: usize,
3160    out: &mut [f32],
3161    pool: Option<&Pool>,
3162) {
3163    // NOTE: double-buffering the dequant against the sgemm (a scoped
3164    // thread driving the pool on tile k+1 while the caller multiplies
3165    // tile k) was tried and LOST ~6%: Accelerate's sgemm is itself
3166    // multithreaded, and the dequant workers just steal its cores.
3167    const TR: usize = 2048;
3168    let b = pre.len();
3169    thread_local! {
3170        static XPANEL: std::cell::RefCell<Vec<f32>> = const { std::cell::RefCell::new(Vec::new()) };
3171        static WTILE: std::cell::RefCell<Vec<f32>> = const { std::cell::RefCell::new(Vec::new()) };
3172    }
3173    XPANEL.with(|xp| {
3174        WTILE.with(|wt| {
3175            let mut xpanel = xp.borrow_mut();
3176            xpanel.clear();
3177            for x in pre {
3178                xpanel.extend_from_slice(x);
3179            }
3180            let mut wtile = wt.borrow_mut();
3181            wtile.resize(TR * cols, 0.0);
3182            let mut r0 = 0usize;
3183            while r0 < rows {
3184                let tr = TR.min(rows - r0);
3185                // Dequant the tile (scale folded) — pool-parallel.
3186                let wt_addr = SendMut(wtile.as_mut_ptr());
3187                let run = |start: usize, end: usize| {
3188                    for r in start..end {
3189                        let row = &q[(r0 + r) * cols..(r0 + r + 1) * cols];
3190                        let s = row_scale[r0 + r];
3191                        // SAFETY: workers cover disjoint r ranges.
3192                        let dst =
3193                            unsafe { std::slice::from_raw_parts_mut(wt_addr.at(r * cols), cols) };
3194                        for (d, &v) in dst.iter_mut().zip(row) {
3195                            *d = (v as i8) as f32 * s;
3196                        }
3197                    }
3198                };
3199                dispatch_rows(pool, tr, &run);
3200                // C[b, tr] (at column r0 of out[b, rows]) = X · Wtileᵀ
3201                unsafe {
3202                    accel_blas::cblas_sgemm(
3203                        101, // RowMajor
3204                        111, // NoTrans A
3205                        112, // Trans B
3206                        b as i32,
3207                        tr as i32,
3208                        cols as i32,
3209                        1.0,
3210                        xpanel.as_ptr(),
3211                        cols as i32,
3212                        wtile.as_ptr(),
3213                        cols as i32,
3214                        0.0,
3215                        out.as_mut_ptr().add(r0),
3216                        rows as i32,
3217                    );
3218                }
3219                r0 += tr;
3220            }
3221        })
3222    });
3223}
3224
3225fn qmatmat(
3226    q: &[u8],
3227    row_scale: &[f32],
3228    pre: &[std::borrow::Cow<'_, [f32]>],
3229    rows: usize,
3230    cols: usize,
3231    out: &mut [f32],
3232    pool: Option<&Pool>,
3233) {
3234    let b = pre.len();
3235    debug_assert_eq!(out.len(), b * rows);
3236    // Big prefill batches ride the AMX (roadmap PR3): the row×batch
3237    // SDOT loop below peaks near the CPU's dot throughput, an order
3238    // below the matrix units. Small tensors and tiny test models stay
3239    // on the exact integer path.
3240    #[cfg(target_os = "macos")]
3241    if b >= 8 && rows * cols >= 500_000 && accel_gemm_enabled() {
3242        qmatmat_accel(q, row_scale, pre, rows, cols, out, pool);
3243        return;
3244    }
3245    #[cfg(target_arch = "aarch64")]
3246    if sdot_enabled() {
3247        let acts: Vec<SplitAct> = pre.iter().map(|x| split_act(x)).collect();
3248        let out_addr = SendMut(out.as_mut_ptr());
3249        // Blocked 2×4 (mobile prefill: no AMX to fall back on — this
3250        // path IS the ARM prefill GEMM off Apple silicon).
3251        let blocked_ok = blocked_enabled();
3252        let use_i8mm = i8mm_enabled();
3253        if blocked_ok {
3254            let run = |start: usize, end: usize| {
3255                let mut o = start;
3256                while o < end {
3257                    if o + 2 <= end {
3258                        let r0 = &q[o * cols..(o + 1) * cols];
3259                        let r1 = &q[(o + 1) * cols..(o + 2) * cols];
3260                        let mut bi = 0usize;
3261                        while bi + 4 <= acts.len() {
3262                            let xs = [
3263                                acts[bi].xq.as_slice(),
3264                                acts[bi + 1].xq.as_slice(),
3265                                acts[bi + 2].xq.as_slice(),
3266                                acts[bi + 3].xq.as_slice(),
3267                            ];
3268                            let d = if use_i8mm {
3269                                unsafe { dot_i8_smmla_2x4(r0, r1, xs) }
3270                            } else {
3271                                unsafe { dot_i8_sdot_2x4(r0, r1, xs) }
3272                            };
3273                            for (r, row) in [r0, r1].into_iter().enumerate() {
3274                                for k in 0..4 {
3275                                    let act = &acts[bi + k];
3276                                    let mut v = d[r][k] as f32 * act.sx;
3277                                    for &(j, xv) in &act.outliers {
3278                                        v += (row[j] as i8) as f32 * xv;
3279                                    }
3280                                    unsafe {
3281                                        *out_addr.at((bi + k) * rows + o + r) = v * row_scale[o + r]
3282                                    };
3283                                }
3284                            }
3285                            bi += 4;
3286                        }
3287                        while bi < acts.len() {
3288                            for (r, row) in [r0, r1].into_iter().enumerate() {
3289                                let v = row_dot_sdot(row, &acts[bi]) * row_scale[o + r];
3290                                unsafe { *out_addr.at(bi * rows + o + r) = v };
3291                            }
3292                            bi += 1;
3293                        }
3294                        o += 2;
3295                    } else {
3296                        let row = &q[o * cols..(o + 1) * cols];
3297                        for (bi, act) in acts.iter().enumerate() {
3298                            let v = row_dot_sdot(row, act) * row_scale[o];
3299                            unsafe { *out_addr.at(bi * rows + o) = v };
3300                        }
3301                        o += 1;
3302                    }
3303                }
3304            };
3305            dispatch_rows(pool, rows, &run);
3306            return;
3307        }
3308        let run = |start: usize, end: usize| {
3309            for o in start..end {
3310                let row = &q[o * cols..(o + 1) * cols];
3311                for (bi, act) in acts.iter().enumerate() {
3312                    let v = row_dot_sdot(row, act) * row_scale[o];
3313                    unsafe { *out_addr.at(bi * rows + o) = v };
3314                }
3315            }
3316        };
3317        dispatch_rows(pool, rows, &run);
3318        return;
3319    }
3320    // x86 A8W8 batch. Non-VNNI parts take the BLOCKED 2×4 kernel
3321    // (roadmap P0: two weight rows' abs() stay in registers across four
3322    // activation streams); VNNI machines keep the per-row bias-trick
3323    // dot, which is already throughput-bound there.
3324    #[cfg(target_arch = "x86_64")]
3325    if avx2_a8w8_enabled() {
3326        let acts: Vec<SplitAct> = pre.iter().map(|x| split_act(x)).collect();
3327        let out_addr = SendMut(out.as_mut_ptr());
3328        // CMF_X86_BLOCKED=0 forces the per-row path (paired in-process
3329        // A/B on noisy shared-vCPU hosts).
3330        let blocked_ok = blocked_enabled();
3331        if !avx512vnni_enabled() && blocked_ok {
3332            let run = |start: usize, end: usize| {
3333                let mut o = start;
3334                while o < end {
3335                    if o + 2 <= end {
3336                        let r0 = &q[o * cols..(o + 1) * cols];
3337                        let r1 = &q[(o + 1) * cols..(o + 2) * cols];
3338                        let mut bi = 0usize;
3339                        while bi + 4 <= acts.len() {
3340                            let xs = [
3341                                acts[bi].xq.as_slice(),
3342                                acts[bi + 1].xq.as_slice(),
3343                                acts[bi + 2].xq.as_slice(),
3344                                acts[bi + 3].xq.as_slice(),
3345                            ];
3346                            let d = unsafe { dot_i8_i8_avx2_2x4(r0, r1, xs) };
3347                            for (r, row) in [r0, r1].into_iter().enumerate() {
3348                                for k in 0..4 {
3349                                    let act = &acts[bi + k];
3350                                    let mut v = d[r][k] as f32 * act.sx;
3351                                    for &(j, xv) in &act.outliers {
3352                                        v += (row[j] as i8) as f32 * xv;
3353                                    }
3354                                    unsafe {
3355                                        *out_addr.at((bi + k) * rows + o + r) = v * row_scale[o + r]
3356                                    };
3357                                }
3358                            }
3359                            bi += 4;
3360                        }
3361                        while bi < acts.len() {
3362                            for (r, row) in [r0, r1].into_iter().enumerate() {
3363                                let v = row_dot_avx2(row, &acts[bi]) * row_scale[o + r];
3364                                unsafe { *out_addr.at(bi * rows + o + r) = v };
3365                            }
3366                            bi += 1;
3367                        }
3368                        o += 2;
3369                    } else {
3370                        let row = &q[o * cols..(o + 1) * cols];
3371                        for (bi, act) in acts.iter().enumerate() {
3372                            let v = row_dot_avx2(row, act) * row_scale[o];
3373                            unsafe { *out_addr.at(bi * rows + o) = v };
3374                        }
3375                        o += 1;
3376                    }
3377                }
3378            };
3379            dispatch_rows(pool, rows, &run);
3380            return;
3381        }
3382        let run = |start: usize, end: usize| {
3383            for o in start..end {
3384                let row = &q[o * cols..(o + 1) * cols];
3385                for (bi, act) in acts.iter().enumerate() {
3386                    let v = row_dot_avx2(row, act) * row_scale[o];
3387                    unsafe { *out_addr.at(bi * rows + o) = v };
3388                }
3389            }
3390        };
3391        dispatch_rows(pool, rows, &run);
3392        return;
3393    }
3394    let out_addr = SendMut(out.as_mut_ptr());
3395    let run = |start: usize, end: usize| {
3396        for o in start..end {
3397            let row = &q[o * cols..(o + 1) * cols];
3398            for (bi, x) in pre.iter().enumerate() {
3399                let mut acc = 0f32;
3400                for j in 0..cols {
3401                    acc += (row[j] as i8) as f32 * x[j];
3402                }
3403                unsafe { *out_addr.at(bi * rows + o) = acc * row_scale[o] };
3404            }
3405        }
3406    };
3407    dispatch_rows(pool, rows, &run);
3408}
3409
3410/// Split rows across pool workers (shared qmatvec pattern). Self-balancing
3411/// — see `Pool::run_rows` for why a static 1/n split is wrong here.
3412fn dispatch_rows(pool: Option<&Pool>, rows: usize, run: &(dyn Fn(usize, usize) + Sync)) {
3413    match pool {
3414        Some(pool) if rows >= 256 => pool.run_rows(rows, run),
3415        _ => run(0, rows),
3416    }
3417}
3418
3419/// Split a q4_block blob into (packed nibbles, f16 group scales).
3420fn q4_split(bytes: &[u8], rows: usize, cols: usize) -> (&[u8], &[u8]) {
3421    let groups = rows * cols / GROUP_SIZE;
3422    bytes.split_at(groups * 16)
3423}
3424
3425/// SIMD unpack for the dominant vbit width B=4 (94% of rows on the
3426/// log2-shape calibration): 16 packed bytes -> 32 centered i8 values.
3427/// vbit packs MSB-first, so the HIGH nibble is the even element
3428/// (opposite of q4_block's lo-first interleave). Centering is u-7.
3429#[inline]
3430fn vbit_fill4(data: &[u8], buf: &mut [u8]) {
3431    #[cfg(target_arch = "aarch64")]
3432    unsafe {
3433        return vbit_fill4_neon(data, buf);
3434    }
3435    #[cfg(target_arch = "x86_64")]
3436    if avx2_enabled() {
3437        return unsafe { vbit_fill4_avx2(data, buf) };
3438    }
3439    #[allow(unreachable_code)]
3440    for (blk, chunk) in buf.chunks_exact_mut(8).enumerate() {
3441        let u = unpack8::<4>(&data[blk * 4..]);
3442        for k in 0..8 {
3443            chunk[k] = (u[k] - 7) as i8 as u8;
3444        }
3445    }
3446}
3447
3448#[cfg(target_arch = "aarch64")]
3449#[target_feature(enable = "neon")]
3450unsafe fn vbit_fill4_neon(data: &[u8], buf: &mut [u8]) {
3451    // SAFETY: buf.len() is a multiple of GROUP_SIZE=32; data holds
3452    // buf.len()/2 packed bytes (validated at load).
3453    unsafe {
3454        use core::arch::aarch64::*;
3455        let n = buf.len();
3456        let mask = vdupq_n_u8(0x0F);
3457        let seven = vdupq_n_s8(7);
3458        let mut g = 0usize;
3459        while g * 32 + 32 <= n {
3460            let b = vld1q_u8(data.as_ptr().add(g * 16));
3461            let hi = vshrq_n_u8::<4>(b);
3462            let lo = vandq_u8(b, mask);
3463            let z0 = vsubq_s8(vreinterpretq_s8_u8(vzip1q_u8(hi, lo)), seven);
3464            let z1 = vsubq_s8(vreinterpretq_s8_u8(vzip2q_u8(hi, lo)), seven);
3465            vst1q_u8(buf.as_mut_ptr().add(g * 32), vreinterpretq_u8_s8(z0));
3466            vst1q_u8(buf.as_mut_ptr().add(g * 32 + 16), vreinterpretq_u8_s8(z1));
3467            g += 1;
3468        }
3469    }
3470}
3471
3472#[cfg(target_arch = "x86_64")]
3473#[target_feature(enable = "avx2")]
3474unsafe fn vbit_fill4_avx2(data: &[u8], buf: &mut [u8]) {
3475    // SAFETY: see vbit_fill4_neon.
3476    unsafe {
3477        use core::arch::x86_64::*;
3478        let n = buf.len();
3479        let mask = _mm_set1_epi8(0x0F);
3480        let seven = _mm256_set1_epi8(7);
3481        let mut g = 0usize;
3482        while g * 32 + 32 <= n {
3483            let b = _mm_loadu_si128(data.as_ptr().add(g * 16) as *const __m128i);
3484            let hi = _mm_and_si128(_mm_srli_epi16::<4>(b), mask);
3485            let lo = _mm_and_si128(b, mask);
3486            let z = _mm256_sub_epi8(
3487                _mm256_set_m128i(_mm_unpackhi_epi8(hi, lo), _mm_unpacklo_epi8(hi, lo)),
3488                seven,
3489            );
3490            _mm256_storeu_si256(buf.as_mut_ptr().add(g * 32) as *mut __m256i, z);
3491            g += 1;
3492        }
3493    }
3494}
3495
3496/// Unpack 8 MSB-first B-bit values from exactly B bytes (fixed shifts —
3497/// no serial bit-buffer, auto-vectorizable). Every 32-value group starts
3498/// byte-aligned (32·B/8 is integral for B∈3..8), so groups decompose
3499/// into 4 such blocks.
3500#[inline(always)]
3501fn unpack8<const B: usize>(data: &[u8]) -> [i32; 8] {
3502    let mut acc = 0u64;
3503    for i in 0..B {
3504        acc = (acc << 8) | data[i] as u64;
3505    }
3506    let mask = (1u64 << B) - 1;
3507    let mut out = [0i32; 8];
3508    for (k, o) in out.iter_mut().enumerate() {
3509        *o = ((acc >> ((7 - k) * B)) & mask) as i32;
3510    }
3511    out
3512}
3513
3514/// Fused vbit matvec straight from the mapped bytes (spec §3, P13
3515/// FIG.3): [u8 bits: rows][f16 scales: rows·cols/32][bit-packed rows,
3516/// MSB-first, byte-padded]. Row data offsets are precomputed at load
3517/// (`vbit_row_offsets`) — the per-call prefix scan was O(rows) pure
3518/// overhead on every matvec.
3519#[allow(clippy::too_many_arguments)]
3520fn vbitmatvec(
3521    bytes: &[u8],
3522    offsets: &[usize],
3523    x: &[f32],
3524    rows: usize,
3525    cols: usize,
3526    out: &mut [f32],
3527    pool: Option<&Pool>,
3528) {
3529    debug_assert_eq!(out.len(), rows);
3530    debug_assert_eq!(offsets.len(), rows + 1);
3531
3532    // SDOT path: unpack the row to centered i8 once, then per-group
3533    // int8 dot against the quantized activations — same A8W8 contract
3534    // as q8 (bounded noise; CMF_SDOT=0 keeps the exact scalar path).
3535    if a8w8_enabled() {
3536        let act = split_act(x);
3537        let out_addr = SendMut(out.as_mut_ptr());
3538        let run = move |start: usize, end: usize| {
3539            vbit_range_a8w8(bytes, offsets, x, &act, rows, cols, out_addr, start, end)
3540        };
3541        dispatch_rows(pool, rows, &run);
3542        return;
3543    }
3544
3545    let out_addr = SendMut(out.as_mut_ptr());
3546    let run = move |start: usize, end: usize| {
3547        vbit_range_f32(bytes, offsets, x, rows, cols, out_addr, start, end)
3548    };
3549    dispatch_rows(pool, rows, &run);
3550}
3551
3552/// One vbit row range via the A8W8 int8 path — kernel body of
3553/// `vbitmatvec`, extracted so multi-matrix jobs can drive it for
3554/// several tensors in one dispatch (b=8 rows go exact f32).
3555#[allow(clippy::too_many_arguments)]
3556fn vbit_range_a8w8(
3557    bytes: &[u8],
3558    offsets: &[usize],
3559    x: &[f32],
3560    act: &SplitAct,
3561    rows: usize,
3562    cols: usize,
3563    out: SendMut,
3564    start: usize,
3565    end: usize,
3566) {
3567    let ng = cols / GROUP_SIZE;
3568    let bits = &bytes[..rows];
3569    let sc_off = rows;
3570    let row_dot = |r: usize| -> f32 {
3571        let b = bits[r] as usize;
3572        let l = (1i32 << (b - 1)) - 1;
3573        let mask = (1u64 << b) - 1;
3574        let data = &bytes[offsets[r]..offsets[r + 1]];
3575        if b == 8 {
3576            // u−L reaches 128 → does not fit i8; exact f32 path.
3577            let (mut acc, mut nbits, mut idx) = (0u64, 0usize, 0usize);
3578            let mut dot = 0f32;
3579            for g in 0..ng {
3580                let so = (r * ng + g) * 2;
3581                let sgf = f16_to_f32(u16::from_le_bytes([
3582                    bytes[sc_off + so],
3583                    bytes[sc_off + so + 1],
3584                ]));
3585                let xg = &x[g * GROUP_SIZE..(g + 1) * GROUP_SIZE];
3586                let mut gd = 0f32;
3587                for &xv in xg.iter() {
3588                    if nbits < 8 {
3589                        acc = (acc << 8) | data[idx] as u64;
3590                        idx += 1;
3591                        nbits += 8;
3592                    }
3593                    let u = ((acc >> (nbits - 8)) & 0xFF) as i32;
3594                    nbits -= 8;
3595                    gd += (u - l) as f32 * xv;
3596                }
3597                dot += gd * sgf;
3598            }
3599            return dot;
3600        }
3601        // Per-worker scratch: this closure runs for every row of the
3602        // tensor (lm_head ≈ 150k rows/token) — a heap allocation per
3603        // row was measurable pure overhead.
3604        thread_local! {
3605            static VBIT_SCRATCH: std::cell::RefCell<Vec<u8>> =
3606                const { std::cell::RefCell::new(Vec::new()) };
3607        }
3608        #[inline(always)]
3609        fn fill<const B: usize>(data: &[u8], l: i32, buf: &mut [u8]) {
3610            for (blk, chunk) in buf.chunks_exact_mut(8).enumerate() {
3611                let u = unpack8::<B>(&data[blk * B..]);
3612                for k in 0..8 {
3613                    chunk[k] = (u[k] - l) as i8 as u8;
3614                }
3615            }
3616        }
3617        let _ = mask;
3618        VBIT_SCRATCH.with(|scratch| {
3619            let mut buf = scratch.borrow_mut();
3620            buf.resize(cols, 0);
3621            match b {
3622                3 => fill::<3>(data, l, &mut buf),
3623                4 => vbit_fill4(data, &mut buf),
3624                5 => fill::<5>(data, l, &mut buf),
3625                6 => fill::<6>(data, l, &mut buf),
3626                _ => unreachable!(),
3627            }
3628            let mut dot = 0f32;
3629            for g in 0..ng {
3630                let so = (r * ng + g) * 2;
3631                let s = f16_to_f32(u16::from_le_bytes([
3632                    bytes[sc_off + so],
3633                    bytes[sc_off + so + 1],
3634                ]));
3635                let d = dot_i8_i8(
3636                    &buf[g * GROUP_SIZE..(g + 1) * GROUP_SIZE],
3637                    &act.xq[g * GROUP_SIZE..(g + 1) * GROUP_SIZE],
3638                ) as f32
3639                    * act.sx;
3640                dot += d * s;
3641            }
3642            for &(j, xv) in &act.outliers {
3643                let so = (r * ng + j / GROUP_SIZE) * 2;
3644                let s = f16_to_f32(u16::from_le_bytes([
3645                    bytes[sc_off + so],
3646                    bytes[sc_off + so + 1],
3647                ]));
3648                // xq is zeroed at outlier slots — add the exact term.
3649                dot += (buf[j] as i8) as f32 * s * xv;
3650            }
3651            dot
3652        })
3653    };
3654    for r in start..end {
3655        // SAFETY: disjoint row ranges per worker.
3656        unsafe { *out.at(r) = row_dot(r) };
3657    }
3658}
3659
3660/// Exact scalar vbit row range (same extraction, non-SDOT path).
3661#[allow(clippy::too_many_arguments)]
3662fn vbit_range_f32(
3663    bytes: &[u8],
3664    offsets: &[usize],
3665    x: &[f32],
3666    rows: usize,
3667    cols: usize,
3668    out: SendMut,
3669    start: usize,
3670    end: usize,
3671) {
3672    let ng = cols / GROUP_SIZE;
3673    let bits = &bytes[..rows];
3674    let sc_off = rows;
3675    // Per-bit-width specialized inner loops: the compiler unrolls the
3676    // constant shifts (the generic bit-buffer loop was branch-bound —
3677    // 5.6 vs 13.2 tok/s q4 on the 0.8B).
3678    #[inline(always)]
3679    fn dot_row<const B: usize>(
3680        data: &[u8],
3681        bytes: &[u8],
3682        sc_off: usize,
3683        r: usize,
3684        ng: usize,
3685        x: &[f32],
3686    ) -> f32 {
3687        let l = ((1i32 << (B - 1)) - 1) as f32;
3688        let gbytes = GROUP_SIZE * B / 8;
3689        let mut dot = 0f32;
3690        for g in 0..ng {
3691            let so = (r * ng + g) * 2;
3692            let s = f16_to_f32(u16::from_le_bytes([
3693                bytes[sc_off + so],
3694                bytes[sc_off + so + 1],
3695            ]));
3696            let xg = &x[g * GROUP_SIZE..(g + 1) * GROUP_SIZE];
3697            let gd0 = &data[g * gbytes..(g + 1) * gbytes];
3698            let mut gd = 0f32;
3699            for blk in 0..GROUP_SIZE / 8 {
3700                let u = unpack8::<B>(&gd0[blk * B..]);
3701                let xb = &xg[blk * 8..blk * 8 + 8];
3702                for k in 0..8 {
3703                    gd += (u[k] as f32 - l) * xb[k];
3704                }
3705            }
3706            dot += gd * s;
3707        }
3708        dot
3709    }
3710    for r in start..end {
3711        let data = &bytes[offsets[r]..offsets[r + 1]];
3712        let v = match bits[r] {
3713            3 => dot_row::<3>(data, bytes, sc_off, r, ng, x),
3714            4 => dot_row::<4>(data, bytes, sc_off, r, ng, x),
3715            5 => dot_row::<5>(data, bytes, sc_off, r, ng, x),
3716            6 => dot_row::<6>(data, bytes, sc_off, r, ng, x),
3717            8 => dot_row::<8>(data, bytes, sc_off, r, ng, x),
3718            b => unreachable!("vbit bit-width {b} (validated at load)"),
3719        };
3720        // SAFETY: disjoint row ranges per worker.
3721        unsafe { *out.at(r) = v };
3722    }
3723}
3724
3725/// Fused two-input vbit matvec: each row is unpacked from the mmap ONCE
3726/// and dotted against BOTH activations (MTP verify / pair prefill used
3727/// to run two full matvecs — double weight traffic and double unpack).
3728/// Per-input math is identical to `vbitmatvec` → same accuracy contract.
3729#[allow(clippy::too_many_arguments)]
3730fn vbitmatvec2(
3731    bytes: &[u8],
3732    offsets: &[usize],
3733    x1: &[f32],
3734    x2: &[f32],
3735    rows: usize,
3736    cols: usize,
3737    o1: &mut [f32],
3738    o2: &mut [f32],
3739    pool: Option<&Pool>,
3740) {
3741    debug_assert_eq!(o1.len(), rows);
3742    debug_assert_eq!(o2.len(), rows);
3743
3744    if a8w8_enabled() {
3745        let a1 = split_act(x1);
3746        let a2 = split_act(x2);
3747        let p1 = SendMut(o1.as_mut_ptr());
3748        let p2 = SendMut(o2.as_mut_ptr());
3749        let run = move |start: usize, end: usize| {
3750            vbit_range2_a8w8(
3751                bytes, offsets, x1, x2, &a1, &a2, rows, cols, p1, p2, start, end,
3752            )
3753        };
3754        dispatch_rows(pool, rows, &run);
3755        return;
3756    }
3757
3758    let p1 = SendMut(o1.as_mut_ptr());
3759    let p2 = SendMut(o2.as_mut_ptr());
3760    let run = move |start: usize, end: usize| {
3761        vbit_range2_f32(bytes, offsets, x1, x2, rows, cols, p1, p2, start, end)
3762    };
3763    dispatch_rows(pool, rows, &run);
3764}
3765
3766/// Two-input vbit row range via the A8W8 int8 path — kernel body of
3767/// `vbitmatvec2`, extracted for pair multi-matrix jobs (b=8 rows go
3768/// exact f32 for both lanes, bits streamed once).
3769#[allow(clippy::too_many_arguments)]
3770fn vbit_range2_a8w8(
3771    bytes: &[u8],
3772    offsets: &[usize],
3773    x1: &[f32],
3774    x2: &[f32],
3775    a1: &SplitAct,
3776    a2: &SplitAct,
3777    rows: usize,
3778    cols: usize,
3779    p1: SendMut,
3780    p2: SendMut,
3781    start: usize,
3782    end: usize,
3783) {
3784    let ng = cols / GROUP_SIZE;
3785    let bits = &bytes[..rows];
3786    let sc_off = rows;
3787    let row_dots = |r: usize| -> (f32, f32) {
3788        let b = bits[r] as usize;
3789        let l = (1i32 << (b - 1)) - 1;
3790        let data = &bytes[offsets[r]..offsets[r + 1]];
3791        if b == 8 {
3792            // u−L reaches 128 → does not fit i8; exact f32 path,
3793            // bits still streamed once for both lanes.
3794            let (mut acc, mut nbits, mut idx) = (0u64, 0usize, 0usize);
3795            let (mut d1, mut d2) = (0f32, 0f32);
3796            for g in 0..ng {
3797                let so = (r * ng + g) * 2;
3798                let sgf = f16_to_f32(u16::from_le_bytes([
3799                    bytes[sc_off + so],
3800                    bytes[sc_off + so + 1],
3801                ]));
3802                let (mut g1, mut g2) = (0f32, 0f32);
3803                for k in 0..GROUP_SIZE {
3804                    if nbits < 8 {
3805                        acc = (acc << 8) | data[idx] as u64;
3806                        idx += 1;
3807                        nbits += 8;
3808                    }
3809                    let u = ((acc >> (nbits - 8)) & 0xFF) as i32;
3810                    nbits -= 8;
3811                    let w = (u - l) as f32;
3812                    g1 += w * x1[g * GROUP_SIZE + k];
3813                    g2 += w * x2[g * GROUP_SIZE + k];
3814                }
3815                d1 += g1 * sgf;
3816                d2 += g2 * sgf;
3817            }
3818            return (d1, d2);
3819        }
3820        thread_local! {
3821            static VBIT_SCRATCH2: std::cell::RefCell<Vec<u8>> =
3822                const { std::cell::RefCell::new(Vec::new()) };
3823        }
3824        #[inline(always)]
3825        fn fill<const B: usize>(data: &[u8], l: i32, buf: &mut [u8]) {
3826            for (blk, chunk) in buf.chunks_exact_mut(8).enumerate() {
3827                let u = unpack8::<B>(&data[blk * B..]);
3828                for k in 0..8 {
3829                    chunk[k] = (u[k] - l) as i8 as u8;
3830                }
3831            }
3832        }
3833        VBIT_SCRATCH2.with(|scratch| {
3834            let mut buf = scratch.borrow_mut();
3835            buf.resize(cols, 0);
3836            match b {
3837                3 => fill::<3>(data, l, &mut buf),
3838                4 => vbit_fill4(data, &mut buf),
3839                5 => fill::<5>(data, l, &mut buf),
3840                6 => fill::<6>(data, l, &mut buf),
3841                _ => unreachable!(),
3842            }
3843            let (mut d1, mut d2) = (0f32, 0f32);
3844            for g in 0..ng {
3845                let so = (r * ng + g) * 2;
3846                let s = f16_to_f32(u16::from_le_bytes([
3847                    bytes[sc_off + so],
3848                    bytes[sc_off + so + 1],
3849                ]));
3850                let wg = &buf[g * GROUP_SIZE..(g + 1) * GROUP_SIZE];
3851                let v1 = dot_i8_i8(wg, &a1.xq[g * GROUP_SIZE..(g + 1) * GROUP_SIZE]) as f32 * a1.sx;
3852                let v2 = dot_i8_i8(wg, &a2.xq[g * GROUP_SIZE..(g + 1) * GROUP_SIZE]) as f32 * a2.sx;
3853                d1 += v1 * s;
3854                d2 += v2 * s;
3855            }
3856            for &(j, xv) in &a1.outliers {
3857                let so = (r * ng + j / GROUP_SIZE) * 2;
3858                let s = f16_to_f32(u16::from_le_bytes([
3859                    bytes[sc_off + so],
3860                    bytes[sc_off + so + 1],
3861                ]));
3862                d1 += (buf[j] as i8) as f32 * s * xv;
3863            }
3864            for &(j, xv) in &a2.outliers {
3865                let so = (r * ng + j / GROUP_SIZE) * 2;
3866                let s = f16_to_f32(u16::from_le_bytes([
3867                    bytes[sc_off + so],
3868                    bytes[sc_off + so + 1],
3869                ]));
3870                d2 += (buf[j] as i8) as f32 * s * xv;
3871            }
3872            (d1, d2)
3873        })
3874    };
3875    for r in start..end {
3876        let (v1, v2) = row_dots(r);
3877        // SAFETY: disjoint row ranges per worker.
3878        unsafe {
3879            *p1.at(r) = v1;
3880            *p2.at(r) = v2;
3881        }
3882    }
3883}
3884
3885/// Two-input exact scalar vbit row range (same extraction) —
3886/// per-bit-width specialized, two accumulators per row; per-lane
3887/// accumulation order matches `vbitmatvec` exactly.
3888#[allow(clippy::too_many_arguments)]
3889fn vbit_range2_f32(
3890    bytes: &[u8],
3891    offsets: &[usize],
3892    x1: &[f32],
3893    x2: &[f32],
3894    rows: usize,
3895    cols: usize,
3896    p1: SendMut,
3897    p2: SendMut,
3898    start: usize,
3899    end: usize,
3900) {
3901    let ng = cols / GROUP_SIZE;
3902    let bits = &bytes[..rows];
3903    let sc_off = rows;
3904    #[inline(always)]
3905    #[allow(clippy::too_many_arguments)]
3906    fn dot_row2<const B: usize>(
3907        data: &[u8],
3908        bytes: &[u8],
3909        sc_off: usize,
3910        r: usize,
3911        ng: usize,
3912        x1: &[f32],
3913        x2: &[f32],
3914    ) -> (f32, f32) {
3915        let l = ((1i32 << (B - 1)) - 1) as f32;
3916        let gbytes = GROUP_SIZE * B / 8;
3917        let (mut d1, mut d2) = (0f32, 0f32);
3918        for g in 0..ng {
3919            let so = (r * ng + g) * 2;
3920            let s = f16_to_f32(u16::from_le_bytes([
3921                bytes[sc_off + so],
3922                bytes[sc_off + so + 1],
3923            ]));
3924            let x1g = &x1[g * GROUP_SIZE..(g + 1) * GROUP_SIZE];
3925            let x2g = &x2[g * GROUP_SIZE..(g + 1) * GROUP_SIZE];
3926            let gd0 = &data[g * gbytes..(g + 1) * gbytes];
3927            let (mut g1, mut g2) = (0f32, 0f32);
3928            for blk in 0..GROUP_SIZE / 8 {
3929                let u = unpack8::<B>(&gd0[blk * B..]);
3930                for k in 0..8 {
3931                    let w = u[k] as f32 - l;
3932                    g1 += w * x1g[blk * 8 + k];
3933                    g2 += w * x2g[blk * 8 + k];
3934                }
3935            }
3936            d1 += g1 * s;
3937            d2 += g2 * s;
3938        }
3939        (d1, d2)
3940    }
3941    for r in start..end {
3942        let data = &bytes[offsets[r]..offsets[r + 1]];
3943        let (v1, v2) = match bits[r] {
3944            3 => dot_row2::<3>(data, bytes, sc_off, r, ng, x1, x2),
3945            4 => dot_row2::<4>(data, bytes, sc_off, r, ng, x1, x2),
3946            5 => dot_row2::<5>(data, bytes, sc_off, r, ng, x1, x2),
3947            6 => dot_row2::<6>(data, bytes, sc_off, r, ng, x1, x2),
3948            8 => dot_row2::<8>(data, bytes, sc_off, r, ng, x1, x2),
3949            b => unreachable!("vbit bit-width {b} (validated at load)"),
3950        };
3951        // SAFETY: disjoint row ranges per worker.
3952        unsafe {
3953            *p1.at(r) = v1;
3954            *p2.at(r) = v2;
3955        }
3956    }
3957}
3958
3959// ───────────────────── q4_tiled kernels (§4.3) ─────────────────────
3960
3961/// One q4_tiled row dot on the A8W8 int8 path: per 32-group the tile
3962/// is ONE sequential read — [f16 scale][16B nibbles] — versus the two
3963/// distant streams of the split layout. Values/order identical to the
3964/// split kernels.
3965#[inline]
3966#[allow(unreachable_code)]
3967fn dot_q4t_row_i8(bytes: &[u8], r: usize, gpr: usize, xq: &[i8]) -> f32 {
3968    #[cfg(target_arch = "aarch64")]
3969    unsafe {
3970        return dot_q4t_row_sdot(bytes, r, gpr, xq);
3971    }
3972    #[cfg(target_arch = "x86_64")]
3973    unsafe {
3974        if vnni_tiles_enabled() {
3975            return dot_q4t_row_vnni(bytes, r, gpr, xq);
3976        }
3977        return dot_q4t_row_avx2(bytes, r, gpr, xq);
3978    }
3979    let mut acc = 0f32;
3980    for gi in 0..gpr {
3981        let tile = &bytes[(r * gpr + gi) * Q4_TILE..(r * gpr + gi + 1) * Q4_TILE];
3982        let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
3983        let mut d = 0i32;
3984        for (k, &b) in tile[2..].iter().enumerate() {
3985            d += ((b & 0x0F) as i32 - 8) * xq[gi * GROUP_SIZE + k * 2] as i32
3986                + (((b >> 4) & 0x0F) as i32 - 8) * xq[gi * GROUP_SIZE + k * 2 + 1] as i32;
3987        }
3988        acc += d as f32 * s;
3989    }
3990    acc
3991}
3992
3993#[cfg(target_arch = "aarch64")]
3994#[target_feature(enable = "neon,dotprod")]
3995unsafe fn dot_q4t_row_sdot(bytes: &[u8], r: usize, gpr: usize, xq: &[i8]) -> f32 {
3996    // SAFETY: callers uphold slice-length contracts (18B tile per group,
3997    // xq.len() == gpr·GROUP_SIZE).
3998    unsafe {
3999        use core::arch::aarch64::*;
4000        use core::arch::asm;
4001        let lomask = vdupq_n_u8(0x0F);
4002        let eight = vdupq_n_s8(8);
4003        let mut acc = 0f32;
4004        for gi in 0..gpr {
4005            let t = bytes.as_ptr().add((r * gpr + gi) * Q4_TILE);
4006            let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
4007            let b = vld1q_u8(t.add(2));
4008            let lo = vandq_u8(b, lomask);
4009            let hi = vshrq_n_u8::<4>(b);
4010            let e0 = vsubq_s8(vreinterpretq_s8_u8(vzip1q_u8(lo, hi)), eight);
4011            let e1 = vsubq_s8(vreinterpretq_s8_u8(vzip2q_u8(lo, hi)), eight);
4012            let x0 = vld1q_s8(xq.as_ptr().add(gi * GROUP_SIZE));
4013            let x1 = vld1q_s8(xq.as_ptr().add(gi * GROUP_SIZE + 16));
4014            let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
4015            asm!(
4016                "sdot {a0:v}.4s, {e0:v}.16b, {x0:v}.16b",
4017                "sdot {a1:v}.4s, {e1:v}.16b, {x1:v}.16b",
4018                a0 = inout(vreg) a0, a1 = inout(vreg) a1,
4019                e0 = in(vreg) e0, x0 = in(vreg) x0, e1 = in(vreg) e1, x1 = in(vreg) x1,
4020                options(pure, nomem, nostack),
4021            );
4022            acc += vaddvq_s32(vaddq_s32(a0, a1)) as f32 * s;
4023        }
4024        acc
4025    }
4026}
4027
4028#[cfg(target_arch = "x86_64")]
4029#[target_feature(enable = "avx2")]
4030unsafe fn dot_q4t_row_avx2(bytes: &[u8], r: usize, gpr: usize, xq: &[i8]) -> f32 {
4031    // SAFETY: see dot_q4t_row_sdot.
4032    unsafe {
4033        use core::arch::x86_64::*;
4034        let lomask = _mm_set1_epi8(0x0F);
4035        let eight = _mm256_set1_epi8(8);
4036        let ones = _mm256_set1_epi16(1);
4037        let mut acc = 0f32;
4038        for gi in 0..gpr {
4039            let t = bytes.as_ptr().add((r * gpr + gi) * Q4_TILE);
4040            let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
4041            let b = _mm_loadu_si128(t.add(2) as *const __m128i);
4042            let lo = _mm_and_si128(b, lomask);
4043            let hi = _mm_and_si128(_mm_srli_epi16::<4>(b), lomask);
4044            let w = _mm256_sub_epi8(
4045                _mm256_set_m128i(_mm_unpackhi_epi8(lo, hi), _mm_unpacklo_epi8(lo, hi)),
4046                eight,
4047            );
4048            let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
4049            let p16 = _mm256_maddubs_epi16(_mm256_abs_epi8(w), _mm256_sign_epi8(x, w));
4050            let d = _mm256_madd_epi16(p16, ones);
4051            let hi128 = _mm256_extracti128_si256::<1>(d);
4052            let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
4053            let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
4054            let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
4055            acc += _mm_cvtsi128_si32(s32) as f32 * s;
4056        }
4057        acc
4058    }
4059}
4060
4061/// VNNI twin of `dot_q4t_row_avx2`: same unpack, `vpdpbusd` replaces
4062/// the maddubs+madd pair (see `dpbusd_hsum` — sums are bit-identical).
4063/// 256-bit VL encoding, so the VEX `vpsignb` stays usable.
4064#[cfg(target_arch = "x86_64")]
4065#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
4066unsafe fn dot_q4t_row_vnni(bytes: &[u8], r: usize, gpr: usize, xq: &[i8]) -> f32 {
4067    // SAFETY: see dot_q4t_row_sdot.
4068    unsafe {
4069        use core::arch::x86_64::*;
4070        let lomask = _mm_set1_epi8(0x0F);
4071        let eight = _mm256_set1_epi8(8);
4072        let mut acc = 0f32;
4073        for gi in 0..gpr {
4074            let t = bytes.as_ptr().add((r * gpr + gi) * Q4_TILE);
4075            let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
4076            let b = _mm_loadu_si128(t.add(2) as *const __m128i);
4077            let lo = _mm_and_si128(b, lomask);
4078            let hi = _mm_and_si128(_mm_srli_epi16::<4>(b), lomask);
4079            let w = _mm256_sub_epi8(
4080                _mm256_set_m128i(_mm_unpackhi_epi8(lo, hi), _mm_unpacklo_epi8(lo, hi)),
4081                eight,
4082            );
4083            let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
4084            let d = dpbusd_hsum(_mm256_abs_epi8(w), _mm256_sign_epi8(x, w));
4085            acc += d as f32 * s;
4086        }
4087        acc
4088    }
4089}
4090
4091/// One q4_tiled row against FOUR activation streams: the nibble unpack
4092/// and abs() happen once per group instead of once per (group,
4093/// activation) — the unpack is the dominant per-element cost of the
4094/// tiled format (roadmap P0 portable blocking, q4t leg).
4095#[cfg(target_arch = "x86_64")]
4096// `fma` is NOT implied by `avx2`: without it LLVM lowers _mm256_fmadd_ps
4097// to a libm call per lane — measured 2x slower than the reduction this
4098// kernel replaces. The runtime gate (`avx2_enabled`) already requires
4099// both features, so declaring it here is safe.
4100#[target_feature(enable = "avx2,fma")]
4101unsafe fn dot_q4t_row_1x4_avx2(bytes: &[u8], r: usize, gpr: usize, xs: [&[i8]; 4]) -> [f32; 4] {
4102    // SAFETY: callers uphold the 18B-tile and xq-length contracts.
4103    unsafe {
4104        use core::arch::x86_64::*;
4105        let lomask = _mm_set1_epi8(0x0F);
4106        let eight = _mm256_set1_epi8(8);
4107        let ones = _mm256_set1_epi16(1);
4108        // One f32 accumulator VECTOR per activation, reduced once at the
4109        // end. Folding each group's i32 lanes to a scalar inside the loop
4110        // costs an extracti128 + three shift/add + a movd — a cross-lane
4111        // dependency chain per (group, activation), 288 of them per row at
4112        // cols=2304. The per-group scale is what forces a float
4113        // accumulator; it does not force a horizontal sum.
4114        //
4115        // The four accumulators are NAMED, not an array: as `[__m256; 4]`
4116        // indexed by a loop variable LLVM keeps them in memory and every
4117        // group pays four 32-byte loads and stores. That alone made this
4118        // kernel 2x SLOWER than the per-group reduction it replaces
4119        // (measured on the EPYC box: 150 s vs 71 s for two 256² steps).
4120        let mut f0 = _mm256_setzero_ps();
4121        let mut f1 = _mm256_setzero_ps();
4122        let mut f2 = _mm256_setzero_ps();
4123        let mut f3 = _mm256_setzero_ps();
4124        for gi in 0..gpr {
4125            let t = bytes.as_ptr().add((r * gpr + gi) * Q4_TILE);
4126            let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
4127            let sv = _mm256_set1_ps(s);
4128            let bb = _mm_loadu_si128(t.add(2) as *const __m128i);
4129            let lo = _mm_and_si128(bb, lomask);
4130            let hi = _mm_and_si128(_mm_srli_epi16::<4>(bb), lomask);
4131            let w = _mm256_sub_epi8(
4132                _mm256_set_m128i(_mm_unpackhi_epi8(lo, hi), _mm_unpacklo_epi8(lo, hi)),
4133                eight,
4134            );
4135            let aw = _mm256_abs_epi8(w);
4136            let off = gi * GROUP_SIZE;
4137            let dot = |xq: &[i8]| {
4138                let x = _mm256_loadu_si256(xq.as_ptr().add(off) as *const __m256i);
4139                let p16 = _mm256_maddubs_epi16(aw, _mm256_sign_epi8(x, w));
4140                _mm256_cvtepi32_ps(_mm256_madd_epi16(p16, ones))
4141            };
4142            f0 = _mm256_fmadd_ps(dot(xs[0]), sv, f0);
4143            f1 = _mm256_fmadd_ps(dot(xs[1]), sv, f1);
4144            f2 = _mm256_fmadd_ps(dot(xs[2]), sv, f2);
4145            f3 = _mm256_fmadd_ps(dot(xs[3]), sv, f3);
4146        }
4147        [
4148            hsum256_ps(f0),
4149            hsum256_ps(f1),
4150            hsum256_ps(f2),
4151            hsum256_ps(f3),
4152        ]
4153    }
4154}
4155
4156/// Horizontal sum of eight f32 lanes — the one cross-lane reduction the
4157/// blocked kernels pay, once per row instead of once per group.
4158#[cfg(target_arch = "x86_64")]
4159#[target_feature(enable = "avx2")]
4160#[inline]
4161unsafe fn hsum256_ps(v: core::arch::x86_64::__m256) -> f32 {
4162    // SAFETY: pure register arithmetic on the caller's vector.
4163    unsafe {
4164        use core::arch::x86_64::*;
4165        let hi = _mm256_extractf128_ps::<1>(v);
4166        let s = _mm_add_ps(_mm256_castps256_ps128(v), hi);
4167        let s = _mm_add_ps(s, _mm_movehl_ps(s, s));
4168        let s = _mm_add_ss(s, _mm_shuffle_ps::<0x55>(s, s));
4169        _mm_cvtss_f32(s)
4170    }
4171}
4172
4173/// VNNI twin of `dot_q4t_row_1x4_avx2` (see `dpbusd_hsum`).
4174#[cfg(target_arch = "x86_64")]
4175#[target_feature(enable = "avx2,fma,avx512f,avx512bw,avx512vl,avx512vnni")]
4176unsafe fn dot_q4t_row_1x4_vnni(bytes: &[u8], r: usize, gpr: usize, xs: [&[i8]; 4]) -> [f32; 4] {
4177    // SAFETY: callers uphold the 18B-tile and xq-length contracts.
4178    unsafe {
4179        use core::arch::x86_64::*;
4180        let lomask = _mm_set1_epi8(0x0F);
4181        let eight = _mm256_set1_epi8(8);
4182        // Same shape as the AVX2 twin: accumulate in f32 vectors and pay
4183        // one cross-lane reduction per row, not per (group, activation).
4184        let mut f0 = _mm256_setzero_ps();
4185        let mut f1 = _mm256_setzero_ps();
4186        let mut f2 = _mm256_setzero_ps();
4187        let mut f3 = _mm256_setzero_ps();
4188        for gi in 0..gpr {
4189            let t = bytes.as_ptr().add((r * gpr + gi) * Q4_TILE);
4190            let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
4191            let sv = _mm256_set1_ps(s);
4192            let bb = _mm_loadu_si128(t.add(2) as *const __m128i);
4193            let lo = _mm_and_si128(bb, lomask);
4194            let hi = _mm_and_si128(_mm_srli_epi16::<4>(bb), lomask);
4195            let w = _mm256_sub_epi8(
4196                _mm256_set_m128i(_mm_unpackhi_epi8(lo, hi), _mm_unpacklo_epi8(lo, hi)),
4197                eight,
4198            );
4199            let aw = _mm256_abs_epi8(w);
4200            let off = gi * GROUP_SIZE;
4201            let dot = |xq: &[i8]| {
4202                let x = _mm256_loadu_si256(xq.as_ptr().add(off) as *const __m256i);
4203                _mm256_cvtepi32_ps(_mm256_dpbusd_epi32(
4204                    _mm256_setzero_si256(),
4205                    aw,
4206                    _mm256_sign_epi8(x, w),
4207                ))
4208            };
4209            f0 = _mm256_fmadd_ps(dot(xs[0]), sv, f0);
4210            f1 = _mm256_fmadd_ps(dot(xs[1]), sv, f1);
4211            f2 = _mm256_fmadd_ps(dot(xs[2]), sv, f2);
4212            f3 = _mm256_fmadd_ps(dot(xs[3]), sv, f3);
4213        }
4214        let acc = [
4215            hsum256_ps(f0),
4216            hsum256_ps(f1),
4217            hsum256_ps(f2),
4218            hsum256_ps(f3),
4219        ];
4220        acc
4221    }
4222}
4223
4224/// ARM twin of `dot_q4t_row_1x4_avx2`: one nibble unpack per group
4225/// serves FOUR activation streams. Per stream the group order and f32
4226/// accumulation match `dot_q4t_row_sdot` exactly — batch == matvec
4227/// bit-for-bit.
4228#[cfg(target_arch = "aarch64")]
4229#[target_feature(enable = "neon,dotprod")]
4230unsafe fn dot_q4t_row_1x4_sdot(bytes: &[u8], r: usize, gpr: usize, xs: [&[i8]; 4]) -> [f32; 4] {
4231    // SAFETY: callers uphold the 18B-tile and xq-length contracts.
4232    unsafe {
4233        use core::arch::aarch64::*;
4234        use core::arch::asm;
4235        let lomask = vdupq_n_u8(0x0F);
4236        let eight = vdupq_n_s8(8);
4237        let mut acc = [0f32; 4];
4238        for gi in 0..gpr {
4239            let t = bytes.as_ptr().add((r * gpr + gi) * Q4_TILE);
4240            let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
4241            let b = vld1q_u8(t.add(2));
4242            let lo = vandq_u8(b, lomask);
4243            let hi = vshrq_n_u8::<4>(b);
4244            let e0 = vsubq_s8(vreinterpretq_s8_u8(vzip1q_u8(lo, hi)), eight);
4245            let e1 = vsubq_s8(vreinterpretq_s8_u8(vzip2q_u8(lo, hi)), eight);
4246            for (k, xq) in xs.iter().enumerate() {
4247                let x0 = vld1q_s8(xq.as_ptr().add(gi * GROUP_SIZE));
4248                let x1 = vld1q_s8(xq.as_ptr().add(gi * GROUP_SIZE + 16));
4249                let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
4250                asm!(
4251                    "sdot {a0:v}.4s, {e0:v}.16b, {x0:v}.16b",
4252                    "sdot {a1:v}.4s, {e1:v}.16b, {x1:v}.16b",
4253                    a0 = inout(vreg) a0, a1 = inout(vreg) a1,
4254                    e0 = in(vreg) e0, x0 = in(vreg) x0, e1 = in(vreg) e1, x1 = in(vreg) x1,
4255                    options(pure, nomem, nostack),
4256                );
4257                acc[k] += vaddvq_s32(vaddq_s32(a0, a1)) as f32 * s;
4258            }
4259        }
4260        acc
4261    }
4262}
4263
4264/// Exact-term correction for A8W8 outliers on a tiled row.
4265#[inline]
4266fn q4t_outlier(bytes: &[u8], r: usize, gpr: usize, j: usize) -> (f32, f32) {
4267    let gi = j / GROUP_SIZE;
4268    let k = j % GROUP_SIZE;
4269    let tile = &bytes[(r * gpr + gi) * Q4_TILE..(r * gpr + gi + 1) * Q4_TILE];
4270    let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
4271    let byte = tile[2 + k / 2];
4272    let nib = if k & 1 == 0 { byte & 0x0F } else { byte >> 4 };
4273    ((nib as i32 - 8) as f32, s)
4274}
4275
4276/// Exact scalar q4_tiled row (CMF_SDOT=0 contract) — same pairwise
4277/// accumulation shape as `q4_range_f32`.
4278#[inline]
4279fn q4t_row_exact(bytes: &[u8], r: usize, gpr: usize, x: &[f32]) -> f32 {
4280    let mut acc = 0f32;
4281    for gi in 0..gpr {
4282        let tile = &bytes[(r * gpr + gi) * Q4_TILE..(r * gpr + gi + 1) * Q4_TILE];
4283        let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
4284        let xg = &x[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
4285        let mut ga = 0f32;
4286        for (k, &b) in tile[2..].iter().enumerate() {
4287            ga += ((b & 0x0F) as f32 - 8.0) * xg[k * 2]
4288                + (((b >> 4) & 0x0F) as f32 - 8.0) * xg[k * 2 + 1];
4289        }
4290        acc += ga * s;
4291    }
4292    acc
4293}
4294
4295/// Split view of a `q4tp` payload. The three planes are resolved once per
4296/// matvec instead of per row — `q4tp_sections` is cheap, but doing it inside
4297/// the row loop would put a division on the hot path for nothing.
4298struct Q4tpView<'a> {
4299    nib: &'a [u8],
4300    params: &'a [u8],
4301    codes: &'a [u8],
4302    stride: usize,
4303    /// q2tp reads the ladder with rung 0 = exact zero.
4304    zero_rung: bool,
4305}
4306
4307impl<'a> Q4tpView<'a> {
4308    fn new(bytes: &'a [u8], rows: usize, cols: usize) -> Self {
4309        let (params_off, codes_off, stride) = q4tp_sections(rows, cols);
4310        Self {
4311            nib: &bytes[..params_off],
4312            params: &bytes[params_off..codes_off],
4313            codes: &bytes[codes_off..],
4314            stride,
4315            zero_rung: false,
4316        }
4317    }
4318
4319    /// The q2tp view: identical params/codes planes, 8 B weight chunks.
4320    fn new_q2(bytes: &'a [u8], rows: usize, cols: usize) -> Self {
4321        let (params_off, codes_off, stride) = q2tp_sections(rows, cols);
4322        Self {
4323            nib: &bytes[..params_off],
4324            params: &bytes[params_off..codes_off],
4325            codes: &bytes[codes_off..],
4326            stride,
4327            zero_rung: true,
4328        }
4329    }
4330
4331    /// Expand row `r`'s per-tile scales into `out` (length `gpr`).
4332    ///
4333    /// Doing this once per row — rather than decoding a 5-bit code inside the
4334    /// tile loop — is what makes the format free at runtime. Random access to
4335    /// a packed 5-bit field costs a division, two bounds checks and a branch;
4336    /// the tile's actual work is two `sdot`s, so per-tile decoding dominated
4337    /// the kernel and cost 5x (measured: 1.4 vs 6.9 tok/s on Nanbeige-3B).
4338    /// Walking the plane sequentially with a bit accumulator is ~3 ops.
4339    /// Eight 5-bit codes are exactly five bytes, so a whole group of
4340    /// eight decodes from one little-endian word at fixed shifts. The
4341    /// bit-accumulator this replaces carried a data-dependent `while
4342    /// have < 5` refill whose branch sat in the innermost loop of every
4343    /// q4tp row; a decode profile put this function above the dot
4344    /// products it feeds. Same bitstream, same codes — just no branch
4345    /// and eight independent extractions.
4346    #[inline]
4347    fn scales_into(&self, r: usize, gpr: usize, out: &mut [f32]) {
4348        let tab = if self.zero_rung {
4349            q2tp_ladder(self.params, r)
4350        } else {
4351            q4tp_ladder(self.params, r)
4352        };
4353        let codes = &self.codes[r * self.stride..(r + 1) * self.stride];
4354        let out = &mut out[..gpr];
4355        let mut chunks = out.chunks_exact_mut(8);
4356        let mut ci = 0usize;
4357        for c in &mut chunks {
4358            let w = u64::from(codes[ci])
4359                | u64::from(codes[ci + 1]) << 8
4360                | u64::from(codes[ci + 2]) << 16
4361                | u64::from(codes[ci + 3]) << 24
4362                | u64::from(codes[ci + 4]) << 32;
4363            for (k, o) in c.iter_mut().enumerate() {
4364                *o = tab[((w >> (5 * k)) & 31) as usize];
4365            }
4366            ci += 5;
4367        }
4368        // Fewer than eight codes left: the shared total accessor, which
4369        // tolerates a 5-bit field whose spill byte is past the stride.
4370        let tail = &codes[ci..];
4371        for (k, o) in chunks.into_remainder().iter_mut().enumerate() {
4372            *o = tab[q4tp_code(tail, k)];
4373        }
4374    }
4375}
4376
4377#[inline]
4378fn dot_q4tp_row_i8(nib: &[u8], r: usize, gpr: usize, xq: &[i8], scales: &[f32]) -> f32 {
4379    #[cfg(target_arch = "aarch64")]
4380    unsafe {
4381        return dot_q4tp_row_sdot(nib, r, gpr, xq, scales);
4382    }
4383    #[cfg(target_arch = "x86_64")]
4384    unsafe {
4385        if vnni_tiles_enabled() {
4386            return dot_q4tp_row_vnni(nib, r, gpr, xq, scales);
4387        }
4388        return dot_q4tp_row_avx2(nib, r, gpr, xq, scales);
4389    }
4390    #[allow(unreachable_code)]
4391    {
4392        let mut acc = 0f32;
4393        for gi in 0..gpr {
4394            let tile = &nib[(r * gpr + gi) * Q4TP_NIB..(r * gpr + gi + 1) * Q4TP_NIB];
4395            let s = scales[gi];
4396            let mut d = 0i32;
4397            for (k, &b) in tile.iter().enumerate() {
4398                d += ((b & 0x0F) as i32 - 8) * xq[gi * GROUP_SIZE + k * 2] as i32
4399                    + (((b >> 4) & 0x0F) as i32 - 8) * xq[gi * GROUP_SIZE + k * 2 + 1] as i32;
4400            }
4401            acc += d as f32 * s;
4402        }
4403        acc
4404    }
4405}
4406
4407/// q4tp twin of `dot_q4t_row_sdot`: identical nibble math, but the tile
4408/// stride is 16 B (no inline scale) and the scale is a ladder lookup.
4409#[cfg(target_arch = "aarch64")]
4410#[target_feature(enable = "neon,dotprod")]
4411unsafe fn dot_q4tp_row_sdot(nib: &[u8], r: usize, gpr: usize, xq: &[i8], scales: &[f32]) -> f32 {
4412    // SAFETY: callers uphold slice-length contracts (16B tile per group,
4413    // xq.len() == gpr·GROUP_SIZE, codes covering gpr 5-bit fields).
4414    unsafe {
4415        use core::arch::aarch64::*;
4416        use core::arch::asm;
4417        let lomask = vdupq_n_u8(0x0F);
4418        let eight = vdupq_n_s8(8);
4419        let mut acc = 0f32;
4420        for gi in 0..gpr {
4421            let t = nib.as_ptr().add((r * gpr + gi) * Q4TP_NIB);
4422            let s = *scales.get_unchecked(gi);
4423            let b = vld1q_u8(t);
4424            let lo = vandq_u8(b, lomask);
4425            let hi = vshrq_n_u8::<4>(b);
4426            let e0 = vsubq_s8(vreinterpretq_s8_u8(vzip1q_u8(lo, hi)), eight);
4427            let e1 = vsubq_s8(vreinterpretq_s8_u8(vzip2q_u8(lo, hi)), eight);
4428            let x0 = vld1q_s8(xq.as_ptr().add(gi * GROUP_SIZE));
4429            let x1 = vld1q_s8(xq.as_ptr().add(gi * GROUP_SIZE + 16));
4430            let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
4431            asm!(
4432                "sdot {a0:v}.4s, {e0:v}.16b, {x0:v}.16b",
4433                "sdot {a1:v}.4s, {e1:v}.16b, {x1:v}.16b",
4434                a0 = inout(vreg) a0, a1 = inout(vreg) a1,
4435                e0 = in(vreg) e0, x0 = in(vreg) x0, e1 = in(vreg) e1, x1 = in(vreg) x1,
4436                options(pure, nomem, nostack),
4437            );
4438            acc += vaddvq_s32(vaddq_s32(a0, a1)) as f32 * s;
4439        }
4440        acc
4441    }
4442}
4443
4444#[cfg(target_arch = "x86_64")]
4445#[target_feature(enable = "avx2")]
4446unsafe fn dot_q4tp_row_avx2(nib: &[u8], r: usize, gpr: usize, xq: &[i8], scales: &[f32]) -> f32 {
4447    // SAFETY: see dot_q4tp_row_sdot.
4448    unsafe {
4449        use core::arch::x86_64::*;
4450        let lomask = _mm_set1_epi8(0x0F);
4451        let eight = _mm256_set1_epi8(8);
4452        let ones = _mm256_set1_epi16(1);
4453        let mut acc = 0f32;
4454        for gi in 0..gpr {
4455            let t = nib.as_ptr().add((r * gpr + gi) * Q4TP_NIB);
4456            let s = *scales.get_unchecked(gi);
4457            let b = _mm_loadu_si128(t as *const __m128i);
4458            let lo = _mm_and_si128(b, lomask);
4459            let hi = _mm_and_si128(_mm_srli_epi16::<4>(b), lomask);
4460            let w = _mm256_sub_epi8(
4461                _mm256_set_m128i(_mm_unpackhi_epi8(lo, hi), _mm_unpacklo_epi8(lo, hi)),
4462                eight,
4463            );
4464            let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
4465            let p16 = _mm256_maddubs_epi16(_mm256_abs_epi8(w), _mm256_sign_epi8(x, w));
4466            let d = _mm256_madd_epi16(p16, ones);
4467            let hi128 = _mm256_extracti128_si256::<1>(d);
4468            let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
4469            let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
4470            let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
4471            acc += _mm_cvtsi128_si32(s32) as f32 * s;
4472        }
4473        acc
4474    }
4475}
4476
4477/// VNNI twin of `dot_q4tp_row_avx2` (see `dot_q4t_row_vnni` for why the
4478/// 256-bit VL encoding is the one to use here).
4479#[cfg(target_arch = "x86_64")]
4480#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
4481unsafe fn dot_q4tp_row_vnni(nib: &[u8], r: usize, gpr: usize, xq: &[i8], scales: &[f32]) -> f32 {
4482    // SAFETY: see dot_q4tp_row_sdot.
4483    unsafe {
4484        use core::arch::x86_64::*;
4485        let lomask = _mm_set1_epi8(0x0F);
4486        let eight = _mm256_set1_epi8(8);
4487        let mut acc = 0f32;
4488        for gi in 0..gpr {
4489            let t = nib.as_ptr().add((r * gpr + gi) * Q4TP_NIB);
4490            let s = *scales.get_unchecked(gi);
4491            let b = _mm_loadu_si128(t as *const __m128i);
4492            let lo = _mm_and_si128(b, lomask);
4493            let hi = _mm_and_si128(_mm_srli_epi16::<4>(b), lomask);
4494            let w = _mm256_sub_epi8(
4495                _mm256_set_m128i(_mm_unpackhi_epi8(lo, hi), _mm_unpacklo_epi8(lo, hi)),
4496                eight,
4497            );
4498            let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
4499            acc += dpbusd_hsum(_mm256_abs_epi8(w), _mm256_sign_epi8(x, w)) as f32 * s;
4500        }
4501        acc
4502    }
4503}
4504
4505/// Exact scalar q4tp row — the `CMF_SDOT=0` contract, same pairwise
4506/// accumulation shape as `q4t_row_exact`.
4507#[inline]
4508fn q4tp_row_exact(nib: &[u8], r: usize, gpr: usize, x: &[f32], scales: &[f32]) -> f32 {
4509    let mut acc = 0f32;
4510    for gi in 0..gpr {
4511        let tile = &nib[(r * gpr + gi) * Q4TP_NIB..(r * gpr + gi + 1) * Q4TP_NIB];
4512        let s = scales[gi];
4513        let xg = &x[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
4514        let mut ga = 0f32;
4515        for (k, &b) in tile.iter().enumerate() {
4516            ga += ((b & 0x0F) as f32 - 8.0) * xg[k * 2]
4517                + (((b >> 4) & 0x0F) as f32 - 8.0) * xg[k * 2 + 1];
4518        }
4519        acc += ga * s;
4520    }
4521    acc
4522}
4523
4524/// Single weight of a q4tp tensor — the a8w8 outlier path, which restores
4525/// activation outliers at full precision after the int8 pass.
4526#[inline]
4527fn q4tp_outlier(nib: &[u8], r: usize, gpr: usize, j: usize, scales: &[f32]) -> (f32, f32) {
4528    let (gi, k) = (j / GROUP_SIZE, j % GROUP_SIZE);
4529    let byte = nib[(r * gpr + gi) * Q4TP_NIB + k / 2];
4530    let n = if k & 1 == 0 { byte & 0x0F } else { byte >> 4 };
4531    ((n as i32 - 8) as f32, scales[gi])
4532}
4533
4534/// Fused q4tp matvec (dispatch mirrors `q4t_matvec`).
4535fn q4tp_matvec(
4536    bytes: &[u8],
4537    x: &[f32],
4538    rows: usize,
4539    cols: usize,
4540    out: &mut [f32],
4541    pool: Option<&Pool>,
4542) {
4543    debug_assert_eq!(out.len(), rows);
4544    let gpr = cols / GROUP_SIZE;
4545    let v = Q4tpView::new(bytes, rows, cols);
4546    let out_addr = SendMut(out.as_mut_ptr());
4547    if a8w8_enabled() {
4548        let act = split_act(x);
4549        let run = |start: usize, end: usize| {
4550            // One scratch row of scales per worker — borrowed, not minted.
4551            with_krow(gpr, |sc| {
4552                for r in start..end {
4553                    v.scales_into(r, gpr, sc);
4554                    let mut acc = dot_q4tp_row_i8(v.nib, r, gpr, &act.xq, sc) * act.sx;
4555                    for &(j, xv) in &act.outliers {
4556                        let (w, s) = q4tp_outlier(v.nib, r, gpr, j, sc);
4557                        acc += w * s * xv;
4558                    }
4559                    // SAFETY: disjoint row ranges per worker.
4560                    unsafe { *out_addr.at(r) = acc };
4561                }
4562            })
4563        };
4564        dispatch_rows(pool, rows, &run);
4565        return;
4566    }
4567    let run = |start: usize, end: usize| {
4568        with_krow(gpr, |sc| {
4569            for r in start..end {
4570                v.scales_into(r, gpr, sc);
4571                // SAFETY: disjoint row ranges per worker.
4572                unsafe { *out_addr.at(r) = q4tp_row_exact(v.nib, r, gpr, x, sc) };
4573            }
4574        })
4575    };
4576    dispatch_rows(pool, rows, &run);
4577}
4578
4579/// Fused two-input q4tp matvec — the SwiGLU gate/up pair. Weights and the
4580/// row ladder are read once and spent on both activation streams.
4581#[allow(clippy::too_many_arguments)]
4582fn q4tp_matvec2(
4583    bytes: &[u8],
4584    x1: &[f32],
4585    x2: &[f32],
4586    rows: usize,
4587    cols: usize,
4588    o1: &mut [f32],
4589    o2: &mut [f32],
4590    pool: Option<&Pool>,
4591) {
4592    let gpr = cols / GROUP_SIZE;
4593    let v = Q4tpView::new(bytes, rows, cols);
4594    let (p1, p2) = (SendMut(o1.as_mut_ptr()), SendMut(o2.as_mut_ptr()));
4595    let run = |start: usize, end: usize| {
4596        let mut sc = vec![0f32; gpr];
4597        for r in start..end {
4598            v.scales_into(r, gpr, &mut sc);
4599            // SAFETY: disjoint row ranges per worker.
4600            unsafe {
4601                *p1.at(r) = q4tp_row_exact(v.nib, r, gpr, x1, &sc);
4602                *p2.at(r) = q4tp_row_exact(v.nib, r, gpr, x2, &sc);
4603            }
4604        }
4605    };
4606    dispatch_rows(pool, rows, &run);
4607}
4608
4609/// One q2tp outlier weight at column `j` of row `r`: the 2-bit code and
4610/// its group scale, mirrored on `q4tp_outlier`.
4611#[inline]
4612fn q2tp_outlier(chunks: &[u8], r: usize, gpr: usize, j: usize, scales: &[f32]) -> (f32, f32) {
4613    let (gi, k) = (j / GROUP_SIZE, j % GROUP_SIZE);
4614    let byte = chunks[(r * gpr + gi) * Q2TP_CHUNK + k / 4];
4615    let c = (byte >> (2 * (k % 4))) & 3;
4616    (c as f32 - 1.5, scales[gi])
4617}
4618
4619/// Integer dot of one q2tp row against pre-quantized activations:
4620/// Σ_g s_g · (Σ c·xq − 1.5·Σ xq). The half-integer grid (c − 1.5)
4621/// becomes exact integer math through the group sums — the same trick
4622/// every a8w8 kernel in this file rides. The codes decode into a
4623/// 32-byte scratch in natural order and the dot itself is the shared
4624/// SDOT primitive; elsewhere a scalar integer loop.
4625#[inline]
4626fn dot_q2tp_row_i8(
4627    chunks: &[u8],
4628    r: usize,
4629    gpr: usize,
4630    xq: &[i8],
4631    gsum: &[i32],
4632    scales: &[f32],
4633) -> f32 {
4634    let mut acc = 0f32;
4635    let base = r * gpr * Q2TP_CHUNK;
4636    #[cfg(not(target_arch = "aarch64"))]
4637    let mut codes = [0i8; GROUP_SIZE];
4638    for gi in 0..gpr {
4639        let ch = &chunks[base + gi * Q2TP_CHUNK..base + (gi + 1) * Q2TP_CHUNK];
4640        let xg = &xq[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
4641        #[cfg(target_arch = "aarch64")]
4642        // NEON: the byte's four 2-bit fields land in four lane vectors
4643        // (shift+mask), vld4 de-interleaves xq to match (xj[k] =
4644        // xq[4k+j]), widening MACs accumulate exactly in i32. A scalar
4645        // decode here cost as much as the dot it fed — the profile put
4646        // it at the top of the whole W2 decode.
4647        let dot = unsafe {
4648            use core::arch::aarch64::*;
4649            let b = vld1_u8(ch.as_ptr());
4650            let three = vdup_n_u8(3);
4651            let c0 = vreinterpret_s8_u8(vand_u8(b, three));
4652            let c1 = vreinterpret_s8_u8(vand_u8(vshr_n_u8(b, 2), three));
4653            let c2 = vreinterpret_s8_u8(vand_u8(vshr_n_u8(b, 4), three));
4654            let c3 = vreinterpret_s8_u8(vand_u8(vshr_n_u8(b, 6), three));
4655            let x4 = vld4_s8(xg.as_ptr());
4656            let mut acc4 = vdupq_n_s32(0);
4657            acc4 = vpadalq_s16(acc4, vmull_s8(c0, x4.0));
4658            acc4 = vpadalq_s16(acc4, vmull_s8(c1, x4.1));
4659            acc4 = vpadalq_s16(acc4, vmull_s8(c2, x4.2));
4660            acc4 = vpadalq_s16(acc4, vmull_s8(c3, x4.3));
4661            vaddvq_s32(acc4)
4662        };
4663        #[cfg(not(target_arch = "aarch64"))]
4664        let dot: i32 = {
4665            for (k, &b) in ch.iter().enumerate() {
4666                codes[k * 4] = (b & 3) as i8;
4667                codes[k * 4 + 1] = ((b >> 2) & 3) as i8;
4668                codes[k * 4 + 2] = ((b >> 4) & 3) as i8;
4669                codes[k * 4 + 3] = ((b >> 6) & 3) as i8;
4670            }
4671            codes
4672                .iter()
4673                .zip(xg)
4674                .map(|(&c, &x)| c as i32 * x as i32)
4675                .sum()
4676        };
4677        acc += scales[gi] * (dot as f32 - 1.5 * gsum[gi] as f32);
4678    }
4679    acc
4680}
4681
4682/// Exact f32 dot of one q2tp row: 2-bit fields LSB-first, (c − 1.5)·s.
4683/// Scalar on purpose — the 2-bit class targets the GPU graph; the CPU
4684/// path exists for parity gates and small-machine fallback.
4685fn q2tp_row_exact(chunks: &[u8], r: usize, gpr: usize, x: &[f32], scales: &[f32]) -> f32 {
4686    let mut acc = 0f32;
4687    for gi in 0..gpr {
4688        let ch = &chunks[(r * gpr + gi) * Q2TP_CHUNK..(r * gpr + gi + 1) * Q2TP_CHUNK];
4689        let s = scales[gi];
4690        let xb = &x[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
4691        let mut g = 0f32;
4692        for (k, &b) in ch.iter().enumerate() {
4693            g += ((b & 3) as f32 - 1.5) * xb[k * 4]
4694                + (((b >> 2) & 3) as f32 - 1.5) * xb[k * 4 + 1]
4695                + (((b >> 4) & 3) as f32 - 1.5) * xb[k * 4 + 2]
4696                + (((b >> 6) & 3) as f32 - 1.5) * xb[k * 4 + 3];
4697        }
4698        acc += s * g;
4699    }
4700    acc
4701}
4702
4703fn q2tp_matvec(
4704    bytes: &[u8],
4705    x: &[f32],
4706    rows: usize,
4707    cols: usize,
4708    out: &mut [f32],
4709    pool: Option<&Pool>,
4710) {
4711    debug_assert_eq!(out.len(), rows);
4712    let gpr = cols / GROUP_SIZE;
4713    let v = Q4tpView::new_q2(bytes, rows, cols);
4714    let out_addr = SendMut(out.as_mut_ptr());
4715    // a8w8 fast path (CMF_SDOT=0 keeps the exact scalar walk): integer
4716    // code dots + group sums, exact outlier correction — the same
4717    // contract as every sibling kernel; measured 2-bit rows were the
4718    // only scalar holdout in the family.
4719    if a8w8_enabled() {
4720        let act = split_act(x);
4721        let gsum = q1_group_sums(&act.xq, gpr);
4722        let (act, gsum) = (&act, &gsum);
4723        let run = move |start: usize, end: usize| {
4724            with_krow(gpr, |sc| {
4725                for r in start..end {
4726                    v.scales_into(r, gpr, sc);
4727                    let mut acc = dot_q2tp_row_i8(v.nib, r, gpr, &act.xq, gsum, sc) * act.sx;
4728                    for &(j, xv) in &act.outliers {
4729                        let (w, s) = q2tp_outlier(v.nib, r, gpr, j, sc);
4730                        acc += w * s * xv;
4731                    }
4732                    // SAFETY: disjoint row ranges per worker.
4733                    unsafe { *out_addr.at(r) = acc };
4734                }
4735            })
4736        };
4737        dispatch_rows(pool, rows, &run);
4738        return;
4739    }
4740    let run = |start: usize, end: usize| {
4741        with_krow(gpr, |sc| {
4742            for r in start..end {
4743                v.scales_into(r, gpr, sc);
4744                // SAFETY: disjoint row ranges per worker.
4745                unsafe { *out_addr.at(r) = q2tp_row_exact(v.nib, r, gpr, x, sc) };
4746            }
4747        })
4748    };
4749    dispatch_rows(pool, rows, &run);
4750}
4751
4752/// Fused two-input q2tp matvec — the SwiGLU gate/up pair.
4753#[allow(clippy::too_many_arguments)]
4754fn q2tp_matvec2(
4755    bytes: &[u8],
4756    x1: &[f32],
4757    x2: &[f32],
4758    rows: usize,
4759    cols: usize,
4760    o1: &mut [f32],
4761    o2: &mut [f32],
4762    pool: Option<&Pool>,
4763) {
4764    let gpr = cols / GROUP_SIZE;
4765    let v = Q4tpView::new_q2(bytes, rows, cols);
4766    let (p1, p2) = (SendMut(o1.as_mut_ptr()), SendMut(o2.as_mut_ptr()));
4767    let run = |start: usize, end: usize| {
4768        let mut sc = vec![0f32; gpr];
4769        for r in start..end {
4770            v.scales_into(r, gpr, &mut sc);
4771            // SAFETY: disjoint row ranges per worker.
4772            unsafe {
4773                *p1.at(r) = q2tp_row_exact(v.nib, r, gpr, x1, &sc);
4774                *p2.at(r) = q2tp_row_exact(v.nib, r, gpr, x2, &sc);
4775            }
4776        }
4777    };
4778    dispatch_rows(pool, rows, &run);
4779}
4780
4781/// Batched q2tp matmat: scalar row kernel over every batch column. CPU
4782/// prefill only — decode rides the graph, so plain and correct beats
4783/// clever here.
4784/// Test doors into the host 2-bit kernels: the stand's heap corruption
4785/// pointed at down-shaped tensors, and the private fns need a way to be
4786/// held to a reference without a model file around them.
4787pub fn q2tp_matvec_for_test(bytes: &[u8], x: &[f32], rows: usize, cols: usize, out: &mut [f32]) {
4788    // The facade IS the reference: encoder oracles hold requant output
4789    // to the exact scalar walk. The production dispatch may take the i8
4790    // fast path, whose error scale is the ACTIVATIONS' — a different
4791    // claim than the encoder correctness these tests pin.
4792    let gpr = cols / GROUP_SIZE;
4793    let v = Q4tpView::new_q2(bytes, rows, cols);
4794    with_krow(gpr, |sc| {
4795        for r in 0..rows {
4796            v.scales_into(r, gpr, sc);
4797            out[r] = q2tp_row_exact(v.nib, r, gpr, x, sc);
4798        }
4799    });
4800}
4801
4802pub fn q2tp_matmat_for_test(
4803    bytes: &[u8],
4804    xs_all: &[f32],
4805    b: usize,
4806    rows: usize,
4807    cols: usize,
4808    out: &mut [f32],
4809) {
4810    q2tp_matmat(bytes, xs_all, b, rows, cols, out, None);
4811}
4812
4813fn q2tp_matmat(
4814    bytes: &[u8],
4815    xs_all: &[f32],
4816    b: usize,
4817    rows: usize,
4818    cols: usize,
4819    out: &mut [f32],
4820    pool: Option<&Pool>,
4821) {
4822    debug_assert_eq!(out.len(), b * rows);
4823    let gpr = cols / GROUP_SIZE;
4824    let v = Q4tpView::new_q2(bytes, rows, cols);
4825    let out_addr = SendMut(out.as_mut_ptr());
4826    let run = |start: usize, end: usize| {
4827        let mut sc = vec![0f32; gpr];
4828        for r in start..end {
4829            v.scales_into(r, gpr, &mut sc);
4830            for bi in 0..b {
4831                let x = &xs_all[bi * cols..(bi + 1) * cols];
4832                // SAFETY: disjoint row ranges per worker.
4833                unsafe { *out_addr.at(bi * rows + r) = q2tp_row_exact(v.nib, r, gpr, x, &sc) };
4834            }
4835        }
4836    };
4837    dispatch_rows(pool, rows, &run);
4838}
4839
4840/// The pre-vectorised shape, kept for A/B (`CMF_Q4TP_V1=1`): the
4841/// horizontal add lands once per group per column instead of once per
4842/// row. Same weights, same activations — only the reduction differs.
4843#[cfg(target_arch = "aarch64")]
4844#[target_feature(enable = "neon,dotprod")]
4845unsafe fn dot_q4tp_row_1x4_sdot_v1(
4846    nib: &[u8],
4847    r: usize,
4848    gpr: usize,
4849    xs: [&[i8]; 4],
4850    scales: &[f32],
4851) -> [f32; 4] {
4852    unsafe {
4853        use core::arch::aarch64::*;
4854        use core::arch::asm;
4855        let lomask = vdupq_n_u8(0x0F);
4856        let eight = vdupq_n_s8(8);
4857        let (mut f0, mut f1, mut f2, mut f3) = (0f32, 0f32, 0f32, 0f32);
4858        for gi in 0..gpr {
4859            let t = nib.as_ptr().add((r * gpr + gi) * Q4TP_NIB);
4860            let s = *scales.get_unchecked(gi);
4861            let bb = vld1q_u8(t);
4862            let lo = vandq_u8(bb, lomask);
4863            let hi = vshrq_n_u8::<4>(bb);
4864            let e0 = vsubq_s8(vreinterpretq_s8_u8(vzip1q_u8(lo, hi)), eight);
4865            let e1 = vsubq_s8(vreinterpretq_s8_u8(vzip2q_u8(lo, hi)), eight);
4866            let mut d = [0f32; 4];
4867            for (k, dk) in d.iter_mut().enumerate() {
4868                let x0 = vld1q_s8(xs[k].as_ptr().add(gi * GROUP_SIZE));
4869                let x1 = vld1q_s8(xs[k].as_ptr().add(gi * GROUP_SIZE + 16));
4870                let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
4871                asm!(
4872                    "sdot {a0:v}.4s, {e0:v}.16b, {x0:v}.16b",
4873                    "sdot {a1:v}.4s, {e1:v}.16b, {x1:v}.16b",
4874                    a0 = inout(vreg) a0, a1 = inout(vreg) a1,
4875                    e0 = in(vreg) e0, x0 = in(vreg) x0, e1 = in(vreg) e1, x1 = in(vreg) x1,
4876                    options(pure, nomem, nostack),
4877                );
4878                *dk = vaddvq_s32(vaddq_s32(a0, a1)) as f32 * s;
4879            }
4880            f0 += d[0];
4881            f1 += d[1];
4882            f2 += d[2];
4883            f3 += d[3];
4884        }
4885        [f0, f1, f2, f3]
4886    }
4887}
4888
4889/// Which q4tp batch kernel to run: 1 = the previous one, 2 = the tuned
4890/// one, 0 = decide from the CPU. An atomic rather than a `OnceLock` so a
4891/// benchmark can alternate the two inside one process, where the machine's
4892/// mood — a shared box drifts ±25% between runs — is the same for both.
4893/// What the two mean is per-architecture: on x86 the blocked AVX-512 path
4894/// against the per-column one, on ARM the two reduction shapes.
4895#[allow(dead_code)]
4896static Q4TP_ALT: std::sync::atomic::AtomicU8 = std::sync::atomic::AtomicU8::new(0);
4897
4898/// Blocking pays on x86 only with 512-bit VNNI. With AVX2 alone, four
4899/// columns sharing an unpack still measured slower than the per-column
4900/// path (23.2 ms against 19.4 on a 48-thread EPYC), because that path
4901/// already dequantizes the row once — so the blocked kernel bought a
4902/// second unpack-free pass at the price of half the vector width.
4903#[cfg(target_arch = "x86_64")]
4904fn q4tp_blocked_x86() -> bool {
4905    match Q4TP_ALT.load(std::sync::atomic::Ordering::Relaxed) {
4906        1 => false,
4907        // A forced ON still asks the CPU. The switch exists so a bench can
4908        // pick a kernel, not so it can promise instructions the machine
4909        // does not have — CI caught that as a SIGILL on a runner without
4910        // AVX-512, where the parity test had turned the path on by hand.
4911        2 => avx512vnni_enabled(),
4912        // Deliberately not cached back into the switch: both gates below
4913        // hold their own `OnceLock`, and latching their answer here would
4914        // make a test's override outlive the test that set it.
4915        _ => blocked_enabled() && avx512vnni_enabled(),
4916    }
4917}
4918
4919/// `CMF_Q4TP_V1=1` picks the old reduction shape (A/B only).
4920#[cfg(target_arch = "aarch64")]
4921#[allow(dead_code)]
4922fn q4tp_v1() -> bool {
4923    match Q4TP_ALT.load(std::sync::atomic::Ordering::Relaxed) {
4924        1 => true,
4925        2 => false,
4926        _ => {
4927            static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
4928            *ON.get_or_init(|| std::env::var("CMF_Q4TP_V1").is_ok_and(|v| v != "0"))
4929        }
4930    }
4931}
4932
4933/// Two weight rows against eight columns. The activation load is the
4934/// same for both rows, so it is paid once for twice the arithmetic, and
4935/// sixteen accumulator chains run where eight did — which is what a kernel
4936/// retiring 0.29 instructions a cycle is short of. Register pressure is
4937/// the limit: sixteen `zmm` accumulators, two weight tiles, one
4938/// activation, of thirty-two.
4939///
4940/// Four rows by four columns spends the same sixteen accumulators the
4941/// other way and measured worse — 1488 GFLOP/s against 1644 — so the
4942/// unpack, which four rows pay twice as often, costs more than the extra
4943/// sharing of one activation load buys.
4944#[cfg(target_arch = "x86_64")]
4945#[target_feature(enable = "avx512f,avx512bw,avx512vnni")]
4946unsafe fn dot_q4tp_2x8_avx512(
4947    nib: &[u8],
4948    r0: usize,
4949    gpr: usize,
4950    xs: [&[i8]; 8],
4951    sc0: &[f32],
4952    sc1: &[f32],
4953) -> [[f32; 8]; 2] {
4954    // SAFETY: as dot_q4tp_row_1x8_avx512, two adjacent rows at once; the
4955    // caller guarantees r0 + 1 < rows and the ISA.
4956    unsafe {
4957        use core::arch::x86_64::*;
4958        let lomask = _mm256_set1_epi8(0x0F);
4959        let eight = _mm256_set1_epi8(8);
4960        let zero = _mm512_setzero_si512();
4961        let mut v0 = [_mm512_setzero_ps(); 8];
4962        let mut v1 = [_mm512_setzero_ps(); 8];
4963        let pairs = gpr / 2;
4964        let unpack = |r: usize, gi: usize| -> (__m512i, __mmask64) {
4965            let t = nib.as_ptr().add((r * gpr + gi) * Q4TP_NIB);
4966            let bb = _mm256_loadu_si256(t as *const __m256i);
4967            let lo = _mm256_and_si256(bb, lomask);
4968            let hi = _mm256_and_si256(_mm256_srli_epi16::<4>(bb), lomask);
4969            let ul = _mm256_sub_epi8(_mm256_unpacklo_epi8(lo, hi), eight);
4970            let uh = _mm256_sub_epi8(_mm256_unpackhi_epi8(lo, hi), eight);
4971            let cat = _mm512_inserti64x4::<1>(_mm512_castsi256_si512(ul), uh);
4972            let w = _mm512_shuffle_i64x2::<0b11_01_10_00>(cat, cat);
4973            (_mm512_abs_epi8(w), _mm512_movepi8_mask(w))
4974        };
4975        for gp in 0..pairs {
4976            let gi = gp * 2;
4977            let (wa0, neg0) = unpack(r0, gi);
4978            let (wa1, neg1) = unpack(r0 + 1, gi);
4979            let off = gi * GROUP_SIZE;
4980            let sv = |sc: &[f32]| {
4981                _mm512_insertf32x8::<1>(
4982                    _mm512_castps256_ps512(_mm256_set1_ps(*sc.get_unchecked(gi))),
4983                    _mm256_set1_ps(*sc.get_unchecked(gi + 1)),
4984                )
4985            };
4986            let s0 = sv(sc0);
4987            let s1 = sv(sc1);
4988            for k in 0..8 {
4989                let xv = _mm512_loadu_si512(xs[k].as_ptr().add(off) as *const __m512i);
4990                let d0 = _mm512_cvtepi32_ps(_mm512_dpbusd_epi32(
4991                    zero,
4992                    wa0,
4993                    _mm512_mask_sub_epi8(xv, neg0, zero, xv),
4994                ));
4995                let d1 = _mm512_cvtepi32_ps(_mm512_dpbusd_epi32(
4996                    zero,
4997                    wa1,
4998                    _mm512_mask_sub_epi8(xv, neg1, zero, xv),
4999                ));
5000                v0[k] = _mm512_fmadd_ps(d0, s0, v0[k]);
5001                v1[k] = _mm512_fmadd_ps(d1, s1, v1[k]);
5002            }
5003        }
5004        let mut acc = [[0f32; 8]; 2];
5005        for k in 0..8 {
5006            acc[0][k] = _mm512_reduce_add_ps(v0[k]);
5007            acc[1][k] = _mm512_reduce_add_ps(v1[k]);
5008        }
5009        if gpr % 2 == 1 {
5010            let off = (gpr - 1) * GROUP_SIZE;
5011            for j in off..off + GROUP_SIZE {
5012                let (w0, sa) = q4tp_outlier(nib, r0, gpr, j, sc0);
5013                let (w1, sb) = q4tp_outlier(nib, r0 + 1, gpr, j, sc1);
5014                for k in 0..8 {
5015                    let x = *xs[k].get_unchecked(j) as f32;
5016                    acc[0][k] += w0 * sa * x;
5017                    acc[1][k] += w1 * sb * x;
5018                }
5019            }
5020        }
5021        acc
5022    }
5023}
5024
5025/// The same, eight columns at a time. One unpack then feeds twice as many
5026/// activation streams, so a wide batch reads the weight tile half as
5027/// often; the price is eight accumulators live at once. Measured 9.0 ->
5028/// 8.3 ms at 9216x2304, b=296 on a 48-thread EPYC 9B45.
5029#[cfg(target_arch = "x86_64")]
5030#[target_feature(enable = "avx512f,avx512bw,avx512vnni")]
5031unsafe fn dot_q4tp_row_1x8_avx512(
5032    nib: &[u8],
5033    r: usize,
5034    gpr: usize,
5035    xs: [&[i8]; 8],
5036    scales: &[f32],
5037) -> [f32; 8] {
5038    // SAFETY: as dot_q4tp_row_1x4_avx2; caller guarantees the ISA.
5039    unsafe {
5040        use core::arch::x86_64::*;
5041        let lomask = _mm256_set1_epi8(0x0F);
5042        let eight = _mm256_set1_epi8(8);
5043        let zero = _mm512_setzero_si512();
5044        let (mut v0, mut v1, mut v2, mut v3) = (
5045            _mm512_setzero_ps(),
5046            _mm512_setzero_ps(),
5047            _mm512_setzero_ps(),
5048            _mm512_setzero_ps(),
5049        );
5050        let (mut v4, mut v5, mut v6, mut v7) = (
5051            _mm512_setzero_ps(),
5052            _mm512_setzero_ps(),
5053            _mm512_setzero_ps(),
5054            _mm512_setzero_ps(),
5055        );
5056        let pairs = gpr / 2;
5057        for gp in 0..pairs {
5058            let gi = gp * 2;
5059            let t = nib.as_ptr().add((r * gpr + gi) * Q4TP_NIB);
5060            let bb = _mm256_loadu_si256(t as *const __m256i);
5061            let lo = _mm256_and_si256(bb, lomask);
5062            let hi = _mm256_and_si256(_mm256_srli_epi16::<4>(bb), lomask);
5063            // `unpack` works per 128-bit lane, so the halves come out as
5064            // [A.lo, B.lo] and [A.hi, B.hi]; the shuffle reorders the four
5065            // 128-bit lanes into the weights' natural order, which is what
5066            // the straight activation load expects.
5067            let ul = _mm256_sub_epi8(_mm256_unpacklo_epi8(lo, hi), eight);
5068            let uh = _mm256_sub_epi8(_mm256_unpackhi_epi8(lo, hi), eight);
5069            let cat = _mm512_inserti64x4::<1>(_mm512_castsi256_si512(ul), uh);
5070            let w = _mm512_shuffle_i64x2::<0b11_01_10_00>(cat, cat);
5071            let wabs = _mm512_abs_epi8(w);
5072            let neg = _mm512_movepi8_mask(w);
5073            let off = gi * GROUP_SIZE;
5074            let sv = _mm512_insertf32x8::<1>(
5075                _mm512_castps256_ps512(_mm256_set1_ps(*scales.get_unchecked(gi))),
5076                _mm256_set1_ps(*scales.get_unchecked(gi + 1)),
5077            );
5078            let dot = |x: &[i8]| -> __m512 {
5079                let xv = _mm512_loadu_si512(x.as_ptr().add(off) as *const __m512i);
5080                let sx = _mm512_mask_sub_epi8(xv, neg, zero, xv);
5081                _mm512_cvtepi32_ps(_mm512_dpbusd_epi32(zero, wabs, sx))
5082            };
5083            v0 = _mm512_fmadd_ps(dot(xs[0]), sv, v0);
5084            v1 = _mm512_fmadd_ps(dot(xs[1]), sv, v1);
5085            v2 = _mm512_fmadd_ps(dot(xs[2]), sv, v2);
5086            v3 = _mm512_fmadd_ps(dot(xs[3]), sv, v3);
5087            v4 = _mm512_fmadd_ps(dot(xs[4]), sv, v4);
5088            v5 = _mm512_fmadd_ps(dot(xs[5]), sv, v5);
5089            v6 = _mm512_fmadd_ps(dot(xs[6]), sv, v6);
5090            v7 = _mm512_fmadd_ps(dot(xs[7]), sv, v7);
5091        }
5092        let mut acc = [
5093            _mm512_reduce_add_ps(v0),
5094            _mm512_reduce_add_ps(v1),
5095            _mm512_reduce_add_ps(v2),
5096            _mm512_reduce_add_ps(v3),
5097            _mm512_reduce_add_ps(v4),
5098            _mm512_reduce_add_ps(v5),
5099            _mm512_reduce_add_ps(v6),
5100            _mm512_reduce_add_ps(v7),
5101        ];
5102        // An odd group count leaves one group over; the narrow kernel
5103        // finishes it rather than the tail being a special case here.
5104        if gpr % 2 == 1 {
5105            let off = (gpr - 1) * GROUP_SIZE;
5106            for j in off..off + GROUP_SIZE {
5107                let (w, s) = q4tp_outlier(nib, r, gpr, j, scales);
5108                let ws = w * s;
5109                for k in 0..8 {
5110                    acc[k] += ws * *xs[k].get_unchecked(j) as f32;
5111                }
5112            }
5113        }
5114        acc
5115    }
5116}
5117
5118/// The same four columns, 512 bits wide. Two groups (64 weights) ride one
5119/// unpack and one `vpdpbusd`, where AVX2 needs two unpacks and four
5120/// `maddubs`/`madd` pairs — about 2.3x fewer instructions for the same
5121/// arithmetic. The two groups carry different scales, so the fma takes a
5122/// vector whose halves hold each group's scale rather than a broadcast.
5123///
5124/// There is no 512-bit `vpsignb`, so the activation's sign is applied by
5125/// negating under a mask taken from the weight's sign bits. That mask is
5126/// per-tile, so it is hoisted out of the column loop and the per-column
5127/// cost stays exactly one instruction, as with `sign_epi8`. Weights of
5128/// zero are not zeroed by the mask trick and do not need to be: their
5129/// magnitude is zero, so the product is.
5130#[cfg(target_arch = "x86_64")]
5131#[target_feature(enable = "avx512f,avx512bw,avx512vnni")]
5132unsafe fn dot_q4tp_row_1x4_avx512(
5133    nib: &[u8],
5134    r: usize,
5135    gpr: usize,
5136    xs: [&[i8]; 4],
5137    scales: &[f32],
5138) -> [f32; 4] {
5139    // SAFETY: as dot_q4tp_row_1x4_avx2; caller guarantees the ISA.
5140    unsafe {
5141        use core::arch::x86_64::*;
5142        let lomask = _mm256_set1_epi8(0x0F);
5143        let eight = _mm256_set1_epi8(8);
5144        let zero = _mm512_setzero_si512();
5145        let (mut v0, mut v1, mut v2, mut v3) = (
5146            _mm512_setzero_ps(),
5147            _mm512_setzero_ps(),
5148            _mm512_setzero_ps(),
5149            _mm512_setzero_ps(),
5150        );
5151        let pairs = gpr / 2;
5152        for gp in 0..pairs {
5153            let gi = gp * 2;
5154            let t = nib.as_ptr().add((r * gpr + gi) * Q4TP_NIB);
5155            let bb = _mm256_loadu_si256(t as *const __m256i);
5156            let lo = _mm256_and_si256(bb, lomask);
5157            let hi = _mm256_and_si256(_mm256_srli_epi16::<4>(bb), lomask);
5158            // `unpack` works per 128-bit lane, so the halves come out as
5159            // [A.lo, B.lo] and [A.hi, B.hi]; the shuffle reorders the four
5160            // 128-bit lanes into the weights' natural order, which is what
5161            // the straight activation load expects.
5162            let ul = _mm256_sub_epi8(_mm256_unpacklo_epi8(lo, hi), eight);
5163            let uh = _mm256_sub_epi8(_mm256_unpackhi_epi8(lo, hi), eight);
5164            let cat = _mm512_inserti64x4::<1>(_mm512_castsi256_si512(ul), uh);
5165            let w = _mm512_shuffle_i64x2::<0b11_01_10_00>(cat, cat);
5166            let wabs = _mm512_abs_epi8(w);
5167            let neg = _mm512_movepi8_mask(w);
5168            let off = gi * GROUP_SIZE;
5169            let sv = _mm512_insertf32x8::<1>(
5170                _mm512_castps256_ps512(_mm256_set1_ps(*scales.get_unchecked(gi))),
5171                _mm256_set1_ps(*scales.get_unchecked(gi + 1)),
5172            );
5173            let dot = |x: &[i8]| -> __m512 {
5174                let xv = _mm512_loadu_si512(x.as_ptr().add(off) as *const __m512i);
5175                let sx = _mm512_mask_sub_epi8(xv, neg, zero, xv);
5176                _mm512_cvtepi32_ps(_mm512_dpbusd_epi32(zero, wabs, sx))
5177            };
5178            v0 = _mm512_fmadd_ps(dot(xs[0]), sv, v0);
5179            v1 = _mm512_fmadd_ps(dot(xs[1]), sv, v1);
5180            v2 = _mm512_fmadd_ps(dot(xs[2]), sv, v2);
5181            v3 = _mm512_fmadd_ps(dot(xs[3]), sv, v3);
5182        }
5183        let mut acc = [
5184            _mm512_reduce_add_ps(v0),
5185            _mm512_reduce_add_ps(v1),
5186            _mm512_reduce_add_ps(v2),
5187            _mm512_reduce_add_ps(v3),
5188        ];
5189        // An odd group count leaves one group over; the narrow kernel
5190        // finishes it rather than the tail being a special case here.
5191        if gpr % 2 == 1 {
5192            let off = (gpr - 1) * GROUP_SIZE;
5193            for j in off..off + GROUP_SIZE {
5194                let (w, s) = q4tp_outlier(nib, r, gpr, j, scales);
5195                let ws = w * s;
5196                for k in 0..4 {
5197                    acc[k] += ws * *xs[k].get_unchecked(j) as f32;
5198                }
5199            }
5200        }
5201        acc
5202    }
5203}
5204
5205/// Four batch columns against one q4tp row: the tile is unpacked ONCE and
5206/// spent on four activation streams, which is where a prefill batch stops
5207/// being weight-bandwidth-bound. Twin of `dot_q4t_row_1x4_sdot`.
5208#[cfg(target_arch = "aarch64")]
5209#[target_feature(enable = "neon,dotprod")]
5210unsafe fn dot_q4tp_row_1x4_sdot(
5211    nib: &[u8],
5212    r: usize,
5213    gpr: usize,
5214    xs: [&[i8]; 4],
5215    scales: &[f32],
5216) -> [f32; 4] {
5217    // SAFETY: see dot_q4tp_row_sdot; every xs[k] is gpr·GROUP_SIZE long.
5218    unsafe {
5219        use core::arch::aarch64::*;
5220        use core::arch::asm;
5221        let lomask = vdupq_n_u8(0x0F);
5222        let eight = vdupq_n_s8(8);
5223        // Named accumulators, NOT an array indexed by a loop variable: the
5224        // latter does not stay in registers (the same defect cost 2x in the
5225        // AVX2 q4t kernel and again in WGSL).
5226        //
5227        // They are VECTORS, and the horizontal add happens once at the end
5228        // instead of once per group per column. `vaddvq` is a cross-lane
5229        // reduction — with 72 groups and four columns the old shape paid
5230        // 288 of them per row, each one a dependency stall the pipeline
5231        // cannot hide, to save four float adds. The group's scale now
5232        // rides an fma into the lane accumulators, so the arithmetic per
5233        // group is one convert and one fma. Summation order changes (the
5234        // lanes carry independent partial sums), which is the same
5235        // round-off class the SDOT path already lives in — the strict
5236        // kernel (`CMF_SDOT=0`, what `cortiq ppl` runs) is unchanged and
5237        // stays the reference.
5238        let (mut v0, mut v1, mut v2, mut v3) = (
5239            vdupq_n_f32(0.0),
5240            vdupq_n_f32(0.0),
5241            vdupq_n_f32(0.0),
5242            vdupq_n_f32(0.0),
5243        );
5244        for gi in 0..gpr {
5245            let t = nib.as_ptr().add((r * gpr + gi) * Q4TP_NIB);
5246            let s = *scales.get_unchecked(gi);
5247            let bb = vld1q_u8(t);
5248            let lo = vandq_u8(bb, lomask);
5249            let hi = vshrq_n_u8::<4>(bb);
5250            let e0 = vsubq_s8(vreinterpretq_s8_u8(vzip1q_u8(lo, hi)), eight);
5251            let e1 = vsubq_s8(vreinterpretq_s8_u8(vzip2q_u8(lo, hi)), eight);
5252            let off = gi * GROUP_SIZE;
5253            let dot4 = |x: &[i8]| -> int32x4_t {
5254                let x0 = vld1q_s8(x.as_ptr().add(off));
5255                let x1 = vld1q_s8(x.as_ptr().add(off + 16));
5256                let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
5257                asm!(
5258                    "sdot {a0:v}.4s, {e0:v}.16b, {x0:v}.16b",
5259                    "sdot {a1:v}.4s, {e1:v}.16b, {x1:v}.16b",
5260                    a0 = inout(vreg) a0, a1 = inout(vreg) a1,
5261                    e0 = in(vreg) e0, x0 = in(vreg) x0, e1 = in(vreg) e1, x1 = in(vreg) x1,
5262                    options(pure, nomem, nostack),
5263                );
5264                vaddq_s32(a0, a1)
5265            };
5266            v0 = vfmaq_n_f32(v0, vcvtq_f32_s32(dot4(xs[0])), s);
5267            v1 = vfmaq_n_f32(v1, vcvtq_f32_s32(dot4(xs[1])), s);
5268            v2 = vfmaq_n_f32(v2, vcvtq_f32_s32(dot4(xs[2])), s);
5269            v3 = vfmaq_n_f32(v3, vcvtq_f32_s32(dot4(xs[3])), s);
5270        }
5271        [
5272            vaddvq_f32(v0),
5273            vaddvq_f32(v1),
5274            vaddvq_f32(v2),
5275            vaddvq_f32(v3),
5276        ]
5277    }
5278}
5279
5280/// Fused q4tp matmat — the same three arms `q4t_matmat` has. Shipping only
5281/// the scalar one made Nanbeige-3B decode at 1.2 tok/s against q4t's 5.9:
5282/// the format was fine, the missing arms were the whole regression.
5283fn q4tp_matmat(
5284    bytes: &[u8],
5285    xs_all: &[f32],
5286    b: usize,
5287    rows: usize,
5288    cols: usize,
5289    out: &mut [f32],
5290    pool: Option<&Pool>,
5291) {
5292    debug_assert_eq!(out.len(), b * rows);
5293    let gpr = cols / GROUP_SIZE;
5294    let v = Q4tpView::new(bytes, rows, cols);
5295
5296    // Wide batches ride the AMX through a dequant-tile sgemm, as in q4t.
5297    #[cfg(target_os = "macos")]
5298    if b >= 8 && rows * cols >= 500_000 && accel_gemm_enabled() {
5299        dequant_matmat_accel(
5300            &|r, dst| {
5301                let mut sc = [0f32; 32];
5302                let mut scv;
5303                let s: &[f32] = if gpr <= 32 {
5304                    v.scales_into(r, gpr, &mut sc);
5305                    &sc[..gpr]
5306                } else {
5307                    scv = vec![0f32; gpr];
5308                    v.scales_into(r, gpr, &mut scv);
5309                    &scv
5310                };
5311                for gi in 0..gpr {
5312                    let tile = &v.nib[(r * gpr + gi) * Q4TP_NIB..(r * gpr + gi + 1) * Q4TP_NIB];
5313                    for (k, &bb) in tile.iter().enumerate() {
5314                        dst[gi * GROUP_SIZE + k * 2] = ((bb & 0x0F) as f32 - 8.0) * s[gi];
5315                        dst[gi * GROUP_SIZE + k * 2 + 1] =
5316                            (((bb >> 4) & 0x0F) as f32 - 8.0) * s[gi];
5317                    }
5318                }
5319            },
5320            xs_all,
5321            b,
5322            rows,
5323            cols,
5324            out,
5325            pool,
5326        );
5327        return;
5328    }
5329
5330    let out_addr = SendMut(out.as_mut_ptr());
5331    if a8w8_enabled() {
5332        let acts: Vec<SplitAct> = (0..b)
5333            .map(|bi| split_act(&xs_all[bi * cols..(bi + 1) * cols]))
5334            .collect();
5335        let acts = &acts;
5336        #[cfg(target_arch = "aarch64")]
5337        let blocked_ok = sdot_enabled() && blocked_enabled();
5338        // x86 gets the same blocking: one tile unpack spent on four
5339        // columns. Without it every column re-decoded the row, which is
5340        // why a 48-core EPYC measured a sixth of an M4's per-core rate.
5341        // The gate is `avx2_enabled`, as in q4t — `sdot_enabled` answers
5342        // for ARM's dotprod and is hard-wired false everywhere else, so
5343        // asking it here left the whole blocked path unreachable on x86.
5344        #[cfg(target_arch = "x86_64")]
5345        let blocked_ok = q4tp_blocked_x86();
5346        #[cfg(not(any(target_arch = "aarch64", target_arch = "x86_64")))]
5347        let blocked_ok = false;
5348        // Columns are swept in panels that fit L2. Without this a
5349        // row-pair walks every activation in the batch — 4.8 MB at
5350        // 512x512 — and does it again for the next pair, so the whole
5351        // batch streams out of the shared cache once per row. Measured
5352        // 800 GB/s of it, flat across batch sizes, which is the signature
5353        // of a loop bound by traffic rather than by arithmetic. A panel of
5354        // 256 columns is 590 KB beside 221 KB of this worker's weights:
5355        // both stay resident and the batch crosses L3 once instead of
5356        // once per row.
5357        let panel_cols: usize = std::env::var("CMF_Q4TP_PANEL")
5358            .ok()
5359            .and_then(|v| v.parse().ok())
5360            .filter(|v| *v > 0)
5361            .unwrap_or(256);
5362        let run = |start: usize, end: usize| {
5363            for abase in (0..acts.len()).step_by(panel_cols) {
5364                let alen = (acts.len() - abase).min(panel_cols);
5365                let mut sc = vec![0f32; gpr];
5366                #[cfg(target_arch = "x86_64")]
5367                let mut r_lo = start;
5368                #[cfg(target_arch = "x86_64")]
5369                if blocked_ok && alen >= 8 {
5370                    let mut sc1 = vec![0f32; gpr];
5371                    while r_lo + 2 <= end {
5372                        v.scales_into(r_lo, gpr, &mut sc);
5373                        v.scales_into(r_lo + 1, gpr, &mut sc1);
5374                        let mut bi = 0usize;
5375                        while bi + 8 <= alen {
5376                            let xs = [
5377                                acts[abase + bi].xq.as_slice(),
5378                                acts[abase + bi + 1].xq.as_slice(),
5379                                acts[abase + bi + 2].xq.as_slice(),
5380                                acts[abase + bi + 3].xq.as_slice(),
5381                                acts[abase + bi + 4].xq.as_slice(),
5382                                acts[abase + bi + 5].xq.as_slice(),
5383                                acts[abase + bi + 6].xq.as_slice(),
5384                                acts[abase + bi + 7].xq.as_slice(),
5385                            ];
5386                            let d = unsafe { dot_q4tp_2x8_avx512(v.nib, r_lo, gpr, xs, &sc, &sc1) };
5387                            for (row, dr, scr) in [(r_lo, &d[0], &sc), (r_lo + 1, &d[1], &sc1)] {
5388                                for k in 0..8 {
5389                                    let act = &acts[abase + bi + k];
5390                                    let mut acc = dr[k] * act.sx;
5391                                    for &(j, xv) in &act.outliers {
5392                                        let (w, s) = q4tp_outlier(v.nib, row, gpr, j, scr);
5393                                        acc += w * s * xv;
5394                                    }
5395                                    // SAFETY: disjoint (bi, r) cells per worker.
5396                                    unsafe { *out_addr.at((abase + bi + k) * rows + row) = acc };
5397                                }
5398                            }
5399                            bi += 8;
5400                        }
5401                        // Columns past the last group of eight, both rows —
5402                        // the same single-row kernel the tail below uses.
5403                        for row in [r_lo, r_lo + 1] {
5404                            let scr: &[f32] = if row == r_lo { &sc } else { &sc1 };
5405                            for b2 in bi..alen {
5406                                let act = &acts[abase + b2];
5407                                let xs4 = [
5408                                    act.xq.as_slice(),
5409                                    act.xq.as_slice(),
5410                                    act.xq.as_slice(),
5411                                    act.xq.as_slice(),
5412                                ];
5413                                let d =
5414                                    unsafe { dot_q4tp_row_1x4_avx512(v.nib, row, gpr, xs4, scr) };
5415                                let mut acc = d[0] * act.sx;
5416                                for &(j, xv) in &act.outliers {
5417                                    let (w, s) = q4tp_outlier(v.nib, row, gpr, j, scr);
5418                                    acc += w * s * xv;
5419                                }
5420                                // SAFETY: disjoint (bi, r) cells per worker.
5421                                unsafe { *out_addr.at((abase + b2) * rows + row) = acc };
5422                            }
5423                        }
5424                        r_lo += 2;
5425                    }
5426                }
5427                #[cfg(target_arch = "x86_64")]
5428                let row_start = r_lo;
5429                #[cfg(not(target_arch = "x86_64"))]
5430                let row_start = start;
5431                for r in row_start..end {
5432                    v.scales_into(r, gpr, &mut sc);
5433                    let mut bi = 0usize;
5434                    #[cfg(target_arch = "x86_64")]
5435                    if blocked_ok {
5436                        while bi + 8 <= alen {
5437                            let xs = [
5438                                acts[abase + bi].xq.as_slice(),
5439                                acts[abase + bi + 1].xq.as_slice(),
5440                                acts[abase + bi + 2].xq.as_slice(),
5441                                acts[abase + bi + 3].xq.as_slice(),
5442                                acts[abase + bi + 4].xq.as_slice(),
5443                                acts[abase + bi + 5].xq.as_slice(),
5444                                acts[abase + bi + 6].xq.as_slice(),
5445                                acts[abase + bi + 7].xq.as_slice(),
5446                            ];
5447                            let d = unsafe { dot_q4tp_row_1x8_avx512(v.nib, r, gpr, xs, &sc) };
5448                            for k in 0..8 {
5449                                let act = &acts[abase + bi + k];
5450                                let mut acc = d[k] * act.sx;
5451                                for &(j, xv) in &act.outliers {
5452                                    let (w, s) = q4tp_outlier(v.nib, r, gpr, j, &sc);
5453                                    acc += w * s * xv;
5454                                }
5455                                // SAFETY: disjoint (bi, r) cells per worker.
5456                                unsafe { *out_addr.at((abase + bi + k) * rows + r) = acc };
5457                            }
5458                            bi += 8;
5459                        }
5460                        while bi + 4 <= alen {
5461                            let xs = [
5462                                acts[abase + bi].xq.as_slice(),
5463                                acts[abase + bi + 1].xq.as_slice(),
5464                                acts[abase + bi + 2].xq.as_slice(),
5465                                acts[abase + bi + 3].xq.as_slice(),
5466                            ];
5467                            let d = unsafe { dot_q4tp_row_1x4_avx512(v.nib, r, gpr, xs, &sc) };
5468                            for k in 0..4 {
5469                                let act = &acts[abase + bi + k];
5470                                let mut acc = d[k] * act.sx;
5471                                for &(j, xv) in &act.outliers {
5472                                    let (w, s) = q4tp_outlier(v.nib, r, gpr, j, &sc);
5473                                    acc += w * s * xv;
5474                                }
5475                                // SAFETY: disjoint (bi, r) cells per worker.
5476                                unsafe { *out_addr.at((abase + bi + k) * rows + r) = acc };
5477                            }
5478                            bi += 4;
5479                        }
5480                    }
5481                    #[cfg(target_arch = "aarch64")]
5482                    if blocked_ok {
5483                        while bi + 4 <= alen {
5484                            let xs = [
5485                                acts[abase + bi].xq.as_slice(),
5486                                acts[abase + bi + 1].xq.as_slice(),
5487                                acts[abase + bi + 2].xq.as_slice(),
5488                                acts[abase + bi + 3].xq.as_slice(),
5489                            ];
5490                            let d = unsafe {
5491                                if q4tp_v1() {
5492                                    dot_q4tp_row_1x4_sdot_v1(v.nib, r, gpr, xs, &sc)
5493                                } else {
5494                                    dot_q4tp_row_1x4_sdot(v.nib, r, gpr, xs, &sc)
5495                                }
5496                            };
5497                            for k in 0..4 {
5498                                let act = &acts[abase + bi + k];
5499                                let mut acc = d[k] * act.sx;
5500                                for &(j, xv) in &act.outliers {
5501                                    let (w, s) = q4tp_outlier(v.nib, r, gpr, j, &sc);
5502                                    acc += w * s * xv;
5503                                }
5504                                // SAFETY: disjoint (bi, r) cells per worker.
5505                                unsafe { *out_addr.at((abase + bi + k) * rows + r) = acc };
5506                            }
5507                            bi += 4;
5508                        }
5509                    }
5510                    let _ = blocked_ok;
5511                    while bi < alen {
5512                        let act = &acts[abase + bi];
5513                        let mut acc = dot_q4tp_row_i8(v.nib, r, gpr, &act.xq, &sc) * act.sx;
5514                        for &(j, xv) in &act.outliers {
5515                            let (w, s) = q4tp_outlier(v.nib, r, gpr, j, &sc);
5516                            acc += w * s * xv;
5517                        }
5518                        // SAFETY: disjoint (bi, r) cells per worker range.
5519                        unsafe { *out_addr.at((abase + bi) * rows + r) = acc };
5520                        bi += 1;
5521                    }
5522                }
5523            }
5524        };
5525        dispatch_rows(pool, rows, &run);
5526        return;
5527    }
5528
5529    let run = |start: usize, end: usize| {
5530        let mut sc = vec![0f32; gpr];
5531        for r in start..end {
5532            v.scales_into(r, gpr, &mut sc);
5533            for bi in 0..b {
5534                let x = &xs_all[bi * cols..(bi + 1) * cols];
5535                // SAFETY: disjoint (bi, r) cells per worker range.
5536                unsafe { *out_addr.at(bi * rows + r) = q4tp_row_exact(v.nib, r, gpr, x, &sc) };
5537            }
5538        }
5539    };
5540    dispatch_rows(pool, rows, &run);
5541}
5542
5543/// Fused q4_tiled matvec (dispatch mirrors `q4matvec`).
5544fn q4t_matvec(
5545    bytes: &[u8],
5546    x: &[f32],
5547    rows: usize,
5548    cols: usize,
5549    out: &mut [f32],
5550    pool: Option<&Pool>,
5551) {
5552    debug_assert_eq!(out.len(), rows);
5553    let gpr = cols / GROUP_SIZE;
5554    let out_addr = SendMut(out.as_mut_ptr());
5555    if a8w8_enabled() {
5556        let act = split_act(x);
5557        let run = move |start: usize, end: usize| {
5558            for r in start..end {
5559                let mut acc = dot_q4t_row_i8(bytes, r, gpr, &act.xq) * act.sx;
5560                for &(j, xv) in &act.outliers {
5561                    let (w, s) = q4t_outlier(bytes, r, gpr, j);
5562                    acc += w * s * xv;
5563                }
5564                // SAFETY: disjoint row ranges per worker.
5565                unsafe { *out_addr.at(r) = acc };
5566            }
5567        };
5568        dispatch_rows(pool, rows, &run);
5569        return;
5570    }
5571    let run = move |start: usize, end: usize| {
5572        for r in start..end {
5573            // SAFETY: disjoint row ranges per worker.
5574            unsafe { *out_addr.at(r) = q4t_row_exact(bytes, r, gpr, x) };
5575        }
5576    };
5577    dispatch_rows(pool, rows, &run);
5578}
5579
5580/// Fused two-input q4_tiled matvec (weights read once per pair).
5581#[allow(clippy::too_many_arguments)]
5582fn q4t_matvec2(
5583    bytes: &[u8],
5584    x1: &[f32],
5585    x2: &[f32],
5586    rows: usize,
5587    cols: usize,
5588    o1: &mut [f32],
5589    o2: &mut [f32],
5590    pool: Option<&Pool>,
5591) {
5592    let gpr = cols / GROUP_SIZE;
5593    let p1 = SendMut(o1.as_mut_ptr());
5594    let p2 = SendMut(o2.as_mut_ptr());
5595    if a8w8_enabled() {
5596        let a1 = split_act(x1);
5597        let a2 = split_act(x2);
5598        let run = move |start: usize, end: usize| {
5599            for r in start..end {
5600                let mut v1 = dot_q4t_row_i8(bytes, r, gpr, &a1.xq) * a1.sx;
5601                let mut v2 = dot_q4t_row_i8(bytes, r, gpr, &a2.xq) * a2.sx;
5602                for &(j, xv) in &a1.outliers {
5603                    let (w, s) = q4t_outlier(bytes, r, gpr, j);
5604                    v1 += w * s * xv;
5605                }
5606                for &(j, xv) in &a2.outliers {
5607                    let (w, s) = q4t_outlier(bytes, r, gpr, j);
5608                    v2 += w * s * xv;
5609                }
5610                // SAFETY: disjoint row ranges per worker.
5611                unsafe {
5612                    *p1.at(r) = v1;
5613                    *p2.at(r) = v2;
5614                }
5615            }
5616        };
5617        dispatch_rows(pool, rows, &run);
5618        return;
5619    }
5620    let run = move |start: usize, end: usize| {
5621        for r in start..end {
5622            // SAFETY: disjoint row ranges per worker.
5623            unsafe {
5624                *p1.at(r) = q4t_row_exact(bytes, r, gpr, x1);
5625                *p2.at(r) = q4t_row_exact(bytes, r, gpr, x2);
5626            }
5627        }
5628    };
5629    dispatch_rows(pool, rows, &run);
5630}
5631
5632/// Batched q4_tiled matmat: each row's tiles stream once per microbatch.
5633#[allow(clippy::too_many_arguments)]
5634/// Prefill GEMM through Accelerate for group-quantized codecs: a
5635/// caller-supplied row dequantizer fills f32 tiles (pool-parallel) and
5636/// each tile rides the AMX with one sgemm — the generic sibling of
5637/// `qmatmat_accel` (q8). Numerics are f32-GEMM (tolerance class);
5638/// decode (b=1) never takes this path.
5639#[cfg(target_os = "macos")]
5640fn dequant_matmat_accel(
5641    dequant_row: &(dyn Fn(usize, &mut [f32]) + Sync),
5642    xs_all: &[f32],
5643    b: usize,
5644    rows: usize,
5645    cols: usize,
5646    out: &mut [f32],
5647    pool: Option<&Pool>,
5648) {
5649    const TR: usize = 2048;
5650    thread_local! {
5651        static WTILE: std::cell::RefCell<Vec<f32>> = const { std::cell::RefCell::new(Vec::new()) };
5652    }
5653    WTILE.with(|wt| {
5654        let mut wtile = wt.borrow_mut();
5655        wtile.resize(TR * cols, 0.0);
5656        let mut r0 = 0usize;
5657        while r0 < rows {
5658            let tr = TR.min(rows - r0);
5659            let wt_addr = SendMut(wtile.as_mut_ptr());
5660            let run = |start: usize, end: usize| {
5661                for r in start..end {
5662                    // SAFETY: workers cover disjoint r ranges.
5663                    let dst = unsafe { std::slice::from_raw_parts_mut(wt_addr.at(r * cols), cols) };
5664                    dequant_row(r0 + r, dst);
5665                }
5666            };
5667            dispatch_rows(pool, tr, &run);
5668            unsafe {
5669                accel_blas::cblas_sgemm(
5670                    101, // RowMajor
5671                    111, // NoTrans A
5672                    112, // Trans B
5673                    b as i32,
5674                    tr as i32,
5675                    cols as i32,
5676                    1.0,
5677                    xs_all.as_ptr(),
5678                    cols as i32,
5679                    wtile.as_ptr(),
5680                    cols as i32,
5681                    0.0,
5682                    out.as_mut_ptr().add(r0),
5683                    rows as i32,
5684                );
5685            }
5686            r0 += tr;
5687        }
5688    });
5689}
5690
5691fn q4t_matmat(
5692    bytes: &[u8],
5693    xs_all: &[f32],
5694    b: usize,
5695    rows: usize,
5696    cols: usize,
5697    out: &mut [f32],
5698    pool: Option<&Pool>,
5699) {
5700    debug_assert_eq!(out.len(), b * rows);
5701    let gpr = cols / GROUP_SIZE;
5702    // Wide batches ride the AMX like q8's qmatmat: on Apple silicon
5703    // the dequant-tile sgemm is an order above the SDOT row loop for
5704    // prefill shapes (imagegen DiT forwards are exactly this).
5705    #[cfg(target_os = "macos")]
5706    if b >= 8 && rows * cols >= 500_000 && accel_gemm_enabled() {
5707        dequant_matmat_accel(
5708            &|r, dst| {
5709                for gi in 0..gpr {
5710                    let tile = &bytes[(r * gpr + gi) * Q4_TILE..(r * gpr + gi + 1) * Q4_TILE];
5711                    let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
5712                    for (k, &bb) in tile[2..].iter().enumerate() {
5713                        dst[gi * GROUP_SIZE + k * 2] = ((bb & 0x0F) as f32 - 8.0) * s;
5714                        dst[gi * GROUP_SIZE + k * 2 + 1] = (((bb >> 4) & 0x0F) as f32 - 8.0) * s;
5715                    }
5716                }
5717            },
5718            xs_all,
5719            b,
5720            rows,
5721            cols,
5722            out,
5723            pool,
5724        );
5725        return;
5726    }
5727    let out_addr = SendMut(out.as_mut_ptr());
5728    if a8w8_enabled() {
5729        let acts: Vec<SplitAct> = (0..b)
5730            .map(|bi| split_act(&xs_all[bi * cols..(bi + 1) * cols]))
5731            .collect();
5732        let acts = &acts;
5733        #[cfg(target_arch = "x86_64")]
5734        let blocked_ok = avx2_enabled() && blocked_enabled();
5735        #[cfg(target_arch = "aarch64")]
5736        let blocked_ok = sdot_enabled() && blocked_enabled();
5737        #[cfg(not(any(target_arch = "x86_64", target_arch = "aarch64")))]
5738        let blocked_ok = false;
5739        let run = move |start: usize, end: usize| {
5740            for r in start..end {
5741                let mut bi = 0usize;
5742                #[cfg(target_arch = "aarch64")]
5743                if blocked_ok {
5744                    while bi + 4 <= acts.len() {
5745                        let xs = [
5746                            acts[bi].xq.as_slice(),
5747                            acts[bi + 1].xq.as_slice(),
5748                            acts[bi + 2].xq.as_slice(),
5749                            acts[bi + 3].xq.as_slice(),
5750                        ];
5751                        let d = unsafe { dot_q4t_row_1x4_sdot(bytes, r, gpr, xs) };
5752                        for k in 0..4 {
5753                            let act = &acts[bi + k];
5754                            let mut acc = d[k] * act.sx;
5755                            for &(j, xv) in &act.outliers {
5756                                let (w, sc) = q4t_outlier(bytes, r, gpr, j);
5757                                acc += w * sc * xv;
5758                            }
5759                            // SAFETY: disjoint (bi, r) cells per worker.
5760                            unsafe { *out_addr.at((bi + k) * rows + r) = acc };
5761                        }
5762                        bi += 4;
5763                    }
5764                }
5765                #[cfg(target_arch = "x86_64")]
5766                if blocked_ok {
5767                    while bi + 4 <= acts.len() {
5768                        let xs = [
5769                            acts[bi].xq.as_slice(),
5770                            acts[bi + 1].xq.as_slice(),
5771                            acts[bi + 2].xq.as_slice(),
5772                            acts[bi + 3].xq.as_slice(),
5773                        ];
5774                        let d = unsafe {
5775                            if vnni_tiles_enabled() {
5776                                dot_q4t_row_1x4_vnni(bytes, r, gpr, xs)
5777                            } else {
5778                                dot_q4t_row_1x4_avx2(bytes, r, gpr, xs)
5779                            }
5780                        };
5781                        for k in 0..4 {
5782                            let act = &acts[bi + k];
5783                            let mut acc = d[k] * act.sx;
5784                            for &(j, xv) in &act.outliers {
5785                                let (w, sc) = q4t_outlier(bytes, r, gpr, j);
5786                                acc += w * sc * xv;
5787                            }
5788                            // SAFETY: disjoint (bi, r) cells per worker.
5789                            unsafe { *out_addr.at((bi + k) * rows + r) = acc };
5790                        }
5791                        bi += 4;
5792                    }
5793                }
5794                let _ = blocked_ok;
5795                while bi < acts.len() {
5796                    let act = &acts[bi];
5797                    let mut acc = dot_q4t_row_i8(bytes, r, gpr, &act.xq) * act.sx;
5798                    for &(j, xv) in &act.outliers {
5799                        let (w, s) = q4t_outlier(bytes, r, gpr, j);
5800                        acc += w * s * xv;
5801                    }
5802                    // SAFETY: disjoint (bi, r) cells per worker range.
5803                    unsafe { *out_addr.at(bi * rows + r) = acc };
5804                    bi += 1;
5805                }
5806            }
5807        };
5808        dispatch_rows(pool, rows, &run);
5809        return;
5810    }
5811    let run = move |start: usize, end: usize| {
5812        for r in start..end {
5813            for bi in 0..b {
5814                let x = &xs_all[bi * cols..(bi + 1) * cols];
5815                // SAFETY: disjoint (bi, r) cells per worker range.
5816                unsafe { *out_addr.at(bi * rows + r) = q4t_row_exact(bytes, r, gpr, x) };
5817            }
5818        }
5819    };
5820    dispatch_rows(pool, rows, &run);
5821}
5822
5823// ── q1 (dtype 12): binary weights, [f16 scale][4B sign bits] per
5824// 32-group tile. The kernel family mirrors q4_tiled: one sequential
5825// stream of 6-byte tiles, per-tile integer dot × scale, exact outlier
5826// correction (A8W8 contract), exact scalar path under CMF_SDOT=0. ──
5827
5828/// Per-32-group sums of the quantized activation — the ±1 identity's
5829/// shared half: `dot = −2·sdot(mask, x) − gsum[g]`, computed ONCE per
5830/// matvec and reused by every row.
5831fn q1_group_sums(xq: &[i8], gpr: usize) -> Vec<i32> {
5832    (0..gpr)
5833        .map(|gi| {
5834            xq[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE]
5835                .iter()
5836                .map(|&v| v as i32)
5837                .sum()
5838        })
5839        .collect()
5840}
5841
5842/// One q1 row via the A8W8 int8 path — mask-SDOT on ARM (no ±1
5843/// expansion at all), scalar bit loop elsewhere (AVX2 queued with the
5844/// x86 pass).
5845#[inline]
5846#[allow(unreachable_code)]
5847/// AVX2 q1 row via the same ±1 identity as the ARM sdot kernel: the
5848/// sign bits expand to a {0, −1} byte mask through shuffle+cmpeq, the
5849/// masked activation sums through maddubs(1, x&mask), and
5850/// `dot = −(2·masked_sum + Σx_group)` — bit-identical integer math.
5851#[cfg(target_arch = "x86_64")]
5852#[target_feature(enable = "avx2")]
5853unsafe fn dot_q1_row_avx2(bytes: &[u8], r: usize, gpr: usize, xq: &[i8], gsum: &[i32]) -> f32 {
5854    // SAFETY: callers uphold the 6B-tile and xq/gsum length contracts.
5855    unsafe {
5856        use core::arch::x86_64::*;
5857        // Byte j of the mask must replicate bits-byte j/8.
5858        let expand = _mm256_setr_epi8(
5859            0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 1, 1, 1, 1, 1, 1, 2, 2, 2, 2, 2, 2, 2, 2, 3, 3, 3, 3, 3,
5860            3, 3, 3,
5861        );
5862        let bitsel = _mm256_setr_epi8(
5863            1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64,
5864            -128, 1, 2, 4, 8, 16, 32, 64, -128,
5865        );
5866        let ones8 = _mm256_set1_epi8(1);
5867        let ones16 = _mm256_set1_epi16(1);
5868        let mut acc = 0f32;
5869        for gi in 0..gpr {
5870            let t = bytes.as_ptr().add((r * gpr + gi) * Q1_TILE);
5871            let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
5872            let bits = u32::from_le_bytes([*t.add(2), *t.add(3), *t.add(4), *t.add(5)]);
5873            let bc = _mm256_shuffle_epi8(_mm256_set1_epi32(bits as i32), expand);
5874            let mask = _mm256_cmpeq_epi8(_mm256_and_si256(bc, bitsel), bitsel);
5875            let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
5876            let sel = _mm256_and_si256(x, mask);
5877            // Σ of selected i8 lanes: maddubs(1u8, sel_i8) pairs → madd.
5878            let p16 = _mm256_maddubs_epi16(ones8, sel);
5879            let d32 = _mm256_madd_epi16(p16, ones16);
5880            let hi128 = _mm256_extracti128_si256::<1>(d32);
5881            let s128 = _mm_add_epi32(_mm256_castsi256_si128(d32), hi128);
5882            let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
5883            let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
5884            let msum = _mm_cvtsi128_si32(s32);
5885            // The and-select keeps x UN-negated (unlike ARM's −1-mask
5886            // sdot): d = Σ_set − Σ_unset = 2·Σ_set − Σ_all.
5887            let d = 2 * msum - gsum[gi];
5888            acc += d as f32 * s;
5889        }
5890        acc
5891    }
5892}
5893
5894/// VNNI twin of `dot_q1_row_avx2`: the masked-select sum goes through
5895/// one `vpdpbusd(1u8, sel)` (see `dpbusd_hsum` — bit-identical).
5896#[cfg(target_arch = "x86_64")]
5897#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
5898unsafe fn dot_q1_row_vnni(bytes: &[u8], r: usize, gpr: usize, xq: &[i8], gsum: &[i32]) -> f32 {
5899    // SAFETY: callers uphold the 6B-tile and xq/gsum length contracts.
5900    unsafe {
5901        use core::arch::x86_64::*;
5902        let expand = _mm256_setr_epi8(
5903            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,
5904            3, 3, 3,
5905        );
5906        let bitsel = _mm256_setr_epi8(
5907            1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64,
5908            -128, 1, 2, 4, 8, 16, 32, 64, -128,
5909        );
5910        let ones8 = _mm256_set1_epi8(1);
5911        let mut acc = 0f32;
5912        for gi in 0..gpr {
5913            let t = bytes.as_ptr().add((r * gpr + gi) * Q1_TILE);
5914            let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
5915            let bits = u32::from_le_bytes([*t.add(2), *t.add(3), *t.add(4), *t.add(5)]);
5916            let bc = _mm256_shuffle_epi8(_mm256_set1_epi32(bits as i32), expand);
5917            let mask = _mm256_cmpeq_epi8(_mm256_and_si256(bc, bitsel), bitsel);
5918            let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
5919            let msum = dpbusd_hsum(ones8, _mm256_and_si256(x, mask));
5920            let d = 2 * msum - gsum[gi];
5921            acc += d as f32 * s;
5922        }
5923        acc
5924    }
5925}
5926
5927/// VNNI twin of `dot_q1_row_1x4_avx2` (see `dpbusd_hsum`).
5928#[cfg(target_arch = "x86_64")]
5929#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
5930unsafe fn dot_q1_row_1x4_vnni(
5931    bytes: &[u8],
5932    r: usize,
5933    gpr: usize,
5934    xs: [&[i8]; 4],
5935    gsums: [&[i32]; 4],
5936) -> [f32; 4] {
5937    // SAFETY: callers uphold the 6B-tile and xq/gsum length contracts.
5938    unsafe {
5939        use core::arch::x86_64::*;
5940        let expand = _mm256_setr_epi8(
5941            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,
5942            3, 3, 3,
5943        );
5944        let bitsel = _mm256_setr_epi8(
5945            1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64,
5946            -128, 1, 2, 4, 8, 16, 32, 64, -128,
5947        );
5948        let ones8 = _mm256_set1_epi8(1);
5949        let mut acc = [0f32; 4];
5950        for gi in 0..gpr {
5951            let t = bytes.as_ptr().add((r * gpr + gi) * Q1_TILE);
5952            let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
5953            let bits = u32::from_le_bytes([*t.add(2), *t.add(3), *t.add(4), *t.add(5)]);
5954            let bc = _mm256_shuffle_epi8(_mm256_set1_epi32(bits as i32), expand);
5955            let mask = _mm256_cmpeq_epi8(_mm256_and_si256(bc, bitsel), bitsel);
5956            for (k, xq) in xs.iter().enumerate() {
5957                let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
5958                let msum = dpbusd_hsum(ones8, _mm256_and_si256(x, mask));
5959                let d = 2 * msum - gsums[k][gi];
5960                acc[k] += d as f32 * s;
5961            }
5962        }
5963        acc
5964    }
5965}
5966
5967/// The blocked 1×4 flavor: the expanded bit mask serves four activation
5968/// streams per group (mask build once, four select+reduce chains).
5969#[cfg(target_arch = "x86_64")]
5970#[target_feature(enable = "avx2")]
5971unsafe fn dot_q1_row_1x4_avx2(
5972    bytes: &[u8],
5973    r: usize,
5974    gpr: usize,
5975    xs: [&[i8]; 4],
5976    gsums: [&[i32]; 4],
5977) -> [f32; 4] {
5978    // SAFETY: callers uphold the 6B-tile and xq/gsum length contracts.
5979    unsafe {
5980        use core::arch::x86_64::*;
5981        let expand = _mm256_setr_epi8(
5982            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,
5983            3, 3, 3,
5984        );
5985        let bitsel = _mm256_setr_epi8(
5986            1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64,
5987            -128, 1, 2, 4, 8, 16, 32, 64, -128,
5988        );
5989        let ones8 = _mm256_set1_epi8(1);
5990        let ones16 = _mm256_set1_epi16(1);
5991        let mut acc = [0f32; 4];
5992        for gi in 0..gpr {
5993            let t = bytes.as_ptr().add((r * gpr + gi) * Q1_TILE);
5994            let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
5995            let bits = u32::from_le_bytes([*t.add(2), *t.add(3), *t.add(4), *t.add(5)]);
5996            let bc = _mm256_shuffle_epi8(_mm256_set1_epi32(bits as i32), expand);
5997            let mask = _mm256_cmpeq_epi8(_mm256_and_si256(bc, bitsel), bitsel);
5998            for (k, xq) in xs.iter().enumerate() {
5999                let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
6000                let sel = _mm256_and_si256(x, mask);
6001                let p16 = _mm256_maddubs_epi16(ones8, sel);
6002                let d32 = _mm256_madd_epi16(p16, ones16);
6003                let hi128 = _mm256_extracti128_si256::<1>(d32);
6004                let s128 = _mm_add_epi32(_mm256_castsi256_si128(d32), hi128);
6005                let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
6006                let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
6007                let msum = _mm_cvtsi128_si32(s32);
6008                let d = 2 * msum - gsums[k][gi];
6009                acc[k] += d as f32 * s;
6010            }
6011        }
6012        acc
6013    }
6014}
6015
6016#[allow(unreachable_code)]
6017fn dot_q1_row_i8(bytes: &[u8], r: usize, gpr: usize, xq: &[i8], gsum: &[i32]) -> f32 {
6018    #[cfg(target_arch = "aarch64")]
6019    unsafe {
6020        return dot_q1_row_sdot(bytes, r, gpr, xq, gsum);
6021    }
6022    #[cfg(target_arch = "x86_64")]
6023    if avx2_enabled() {
6024        unsafe {
6025            if vnni_tiles_enabled() {
6026                return dot_q1_row_vnni(bytes, r, gpr, xq, gsum);
6027            }
6028            return dot_q1_row_avx2(bytes, r, gpr, xq, gsum);
6029        }
6030    }
6031    let _ = gsum;
6032    let mut acc = 0f32;
6033    for gi in 0..gpr {
6034        let tile = &bytes[(r * gpr + gi) * Q1_TILE..(r * gpr + gi + 1) * Q1_TILE];
6035        let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
6036        let mut d = 0i32;
6037        for (j, &b) in tile[2..].iter().enumerate() {
6038            for k in 0..8 {
6039                let w = ((b >> k) & 1) as i32 * 2 - 1;
6040                d += w * xq[gi * GROUP_SIZE + j * 8 + k] as i32;
6041            }
6042        }
6043        acc += d as f32 * s;
6044    }
6045    acc
6046}
6047
6048/// SDOT q1 row via the ±1 identity: the vtst mask (0xFF where the bit
6049/// is set, i.e. −1 as i8) feeds `sdot` DIRECTLY — no expansion to ±1
6050/// lanes at all — and `dot = −(2·sdot(mask, x) + Σx_group)`, with the
6051/// per-group activation sums shared across every row of the matvec.
6052/// Four tiles (128 weights) per iteration: integer dots reduce through
6053/// a vpaddq tree into ONE i32x4 that meets its four scales in a single
6054/// fused f32 multiply-add. Integer math throughout — bit-identical to
6055/// the scalar ±1 reference.
6056#[cfg(target_arch = "aarch64")]
6057#[target_feature(enable = "neon,dotprod")]
6058unsafe fn dot_q1_row_sdot(bytes: &[u8], r: usize, gpr: usize, xq: &[i8], gsum: &[i32]) -> f32 {
6059    // SAFETY: callers uphold slice-length contracts (6B tile per group,
6060    // xq.len() == gpr·GROUP_SIZE, gsum.len() == gpr).
6061    unsafe {
6062        use core::arch::aarch64::*;
6063        use core::arch::asm;
6064        const MASKS: [u8; 16] = [1, 2, 4, 8, 16, 32, 64, 128, 1, 2, 4, 8, 16, 32, 64, 128];
6065        let m = vld1q_u8(MASKS.as_ptr());
6066        // One tile's −Σ_set(x) as an UNREDUCED i32x4 (two mask-sdots).
6067        macro_rules! tile_dot {
6068            ($t:expr, $x:expr) => {{
6069                let v0 = vcombine_u8(vdup_n_u8(*$t.add(2)), vdup_n_u8(*$t.add(3)));
6070                let v1 = vcombine_u8(vdup_n_u8(*$t.add(4)), vdup_n_u8(*$t.add(5)));
6071                let w0 = vreinterpretq_s8_u8(vtstq_u8(v0, m));
6072                let w1 = vreinterpretq_s8_u8(vtstq_u8(v1, m));
6073                let x0 = vld1q_s8($x);
6074                let x1 = vld1q_s8($x.add(16));
6075                let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
6076                asm!(
6077                    "sdot {a0:v}.4s, {w0:v}.16b, {x0:v}.16b",
6078                    "sdot {a1:v}.4s, {w1:v}.16b, {x1:v}.16b",
6079                    a0 = inout(vreg) a0, a1 = inout(vreg) a1,
6080                    w0 = in(vreg) w0, x0 = in(vreg) x0, w1 = in(vreg) w1, x1 = in(vreg) x1,
6081                    options(pure, nomem, nostack),
6082                );
6083                vaddq_s32(a0, a1)
6084            }};
6085        }
6086        // TBL unpack over PAIR loads: one vld1q covers two 6B tiles
6087        // ([s s b b b b][s s b b b b] + 4B slack), TBL replicates each
6088        // bit-byte across 8 lanes for vtst, and the four scales gather
6089        // through tbl2 into one fcvtl — the 16 ld1r broadcast loads and
6090        // 4 branchy software f16 conversions per 128 weights (the
6091        // measured load-port wall of this kernel) become 2 vector
6092        // loads + 9 table lookups. Integer math order is unchanged —
6093        // bit-identical results (FCVTL is exact on every f16).
6094        const IW00: [u8; 16] = [2, 2, 2, 2, 2, 2, 2, 2, 3, 3, 3, 3, 3, 3, 3, 3];
6095        const IW01: [u8; 16] = [4, 4, 4, 4, 4, 4, 4, 4, 5, 5, 5, 5, 5, 5, 5, 5];
6096        const IW10: [u8; 16] = [8, 8, 8, 8, 8, 8, 8, 8, 9, 9, 9, 9, 9, 9, 9, 9];
6097        const IW11: [u8; 16] = [
6098            10, 10, 10, 10, 10, 10, 10, 10, 11, 11, 11, 11, 11, 11, 11, 11,
6099        ];
6100        const ISC: [u8; 8] = [0, 1, 6, 7, 16, 17, 22, 23];
6101        let (iw00, iw01) = (vld1q_u8(IW00.as_ptr()), vld1q_u8(IW01.as_ptr()));
6102        let (iw10, iw11) = (vld1q_u8(IW10.as_ptr()), vld1q_u8(IW11.as_ptr()));
6103        let isc = vld1_u8(ISC.as_ptr());
6104        // One tile's −Σ_set(x) from a TBL-unpacked pair load.
6105        macro_rules! tile_dot_tbl {
6106            ($ld:expr, $i0:expr, $i1:expr, $x:expr) => {{
6107                let w0 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8($ld, $i0), m));
6108                let w1 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8($ld, $i1), m));
6109                let x0 = vld1q_s8($x);
6110                let x1 = vld1q_s8($x.add(16));
6111                let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
6112                asm!(
6113                    "sdot {a0:v}.4s, {w0:v}.16b, {x0:v}.16b",
6114                    "sdot {a1:v}.4s, {w1:v}.16b, {x1:v}.16b",
6115                    a0 = inout(vreg) a0, a1 = inout(vreg) a1,
6116                    w0 = in(vreg) w0, x0 = in(vreg) x0, w1 = in(vreg) w1, x1 = in(vreg) x1,
6117                    options(pure, nomem, nostack),
6118                );
6119                vaddq_s32(a0, a1)
6120            }};
6121        }
6122        let base = bytes.as_ptr().add(r * gpr * Q1_TILE);
6123        let row_base = r * gpr * Q1_TILE;
6124        let abs_end = bytes.len();
6125        let xp = xq.as_ptr();
6126        let gp = gsum.as_ptr();
6127        let mut accv = vdupq_n_f32(0.0);
6128        let mut gi = 0;
6129        // The second pair load reads 4B past tile gi+3 — stay inside
6130        // the payload slice (only the file's final tiles fall back).
6131        while gi + 4 <= gpr && row_base + (gi + 4) * Q1_TILE + 4 <= abs_end {
6132            let t0 = base.add(gi * Q1_TILE);
6133            let ld_a = vld1q_u8(t0);
6134            let ld_b = vld1q_u8(t0.add(2 * Q1_TILE));
6135            let d0 = tile_dot_tbl!(ld_a, iw00, iw01, xp.add(gi * GROUP_SIZE));
6136            let d1 = tile_dot_tbl!(ld_a, iw10, iw11, xp.add((gi + 1) * GROUP_SIZE));
6137            let d2 = tile_dot_tbl!(ld_b, iw00, iw01, xp.add((gi + 2) * GROUP_SIZE));
6138            let d3 = tile_dot_tbl!(ld_b, iw10, iw11, xp.add((gi + 3) * GROUP_SIZE));
6139            // [−Σ0, −Σ1, −Σ2, −Σ3] → dots = −(2·Σset_neg + gsum)
6140            let neg = vpaddq_s32(vpaddq_s32(d0, d1), vpaddq_s32(d2, d3));
6141            let g = vld1q_s32(gp.add(gi));
6142            let dots = vnegq_s32(vaddq_s32(vshlq_n_s32::<1>(neg), g));
6143            let sc16 = vqtbl2_u8(uint8x16x2_t(ld_a, ld_b), isc);
6144            let scf: float32x4_t;
6145            asm!(
6146                "fcvtl {o:v}.4s, {i:v}.4h",
6147                o = out(vreg) scf, i = in(vreg) sc16,
6148                options(pure, nomem, nostack),
6149            );
6150            accv = vfmaq_f32(accv, vcvtq_f32_s32(dots), scf);
6151            gi += 4;
6152        }
6153        let mut acc = vaddvq_f32(accv);
6154        while gi < gpr {
6155            let t = base.add(gi * Q1_TILE);
6156            let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
6157            let d = vaddvq_s32(tile_dot!(t, xp.add(gi * GROUP_SIZE)));
6158            acc += (-(2 * d + *gp.add(gi))) as f32 * s;
6159            gi += 1;
6160        }
6161        acc
6162    }
6163}
6164
6165/// Blocked q1 1×4: one TBL unpack of the tile pair serves FOUR
6166/// activation streams (prefill amortization — the same idea as the
6167/// AVX2 twin; per stream the group order, fma order and tail match the
6168/// single-row kernel exactly, so batch == matvec bit-for-bit).
6169#[cfg(target_arch = "aarch64")]
6170#[target_feature(enable = "neon,dotprod")]
6171unsafe fn dot_q1_row_1x4_sdot(
6172    bytes: &[u8],
6173    r: usize,
6174    gpr: usize,
6175    xs: [&[i8]; 4],
6176    gs: [&[i32]; 4],
6177) -> [f32; 4] {
6178    // SAFETY: same slice-length contracts as `dot_q1_row_sdot`, ×4.
6179    unsafe {
6180        use core::arch::aarch64::*;
6181        use core::arch::asm;
6182        const MASKS: [u8; 16] = [1, 2, 4, 8, 16, 32, 64, 128, 1, 2, 4, 8, 16, 32, 64, 128];
6183        const IW00: [u8; 16] = [2, 2, 2, 2, 2, 2, 2, 2, 3, 3, 3, 3, 3, 3, 3, 3];
6184        const IW01: [u8; 16] = [4, 4, 4, 4, 4, 4, 4, 4, 5, 5, 5, 5, 5, 5, 5, 5];
6185        const IW10: [u8; 16] = [8, 8, 8, 8, 8, 8, 8, 8, 9, 9, 9, 9, 9, 9, 9, 9];
6186        const IW11: [u8; 16] = [
6187            10, 10, 10, 10, 10, 10, 10, 10, 11, 11, 11, 11, 11, 11, 11, 11,
6188        ];
6189        const ISC: [u8; 8] = [0, 1, 6, 7, 16, 17, 22, 23];
6190        let m = vld1q_u8(MASKS.as_ptr());
6191        let (iw00, iw01) = (vld1q_u8(IW00.as_ptr()), vld1q_u8(IW01.as_ptr()));
6192        let (iw10, iw11) = (vld1q_u8(IW10.as_ptr()), vld1q_u8(IW11.as_ptr()));
6193        let isc = vld1_u8(ISC.as_ptr());
6194        macro_rules! sdot2 {
6195            ($w0:expr, $w1:expr, $x:expr) => {{
6196                let x0 = vld1q_s8($x);
6197                let x1 = vld1q_s8($x.add(16));
6198                let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
6199                asm!(
6200                    "sdot {a0:v}.4s, {w0:v}.16b, {x0:v}.16b",
6201                    "sdot {a1:v}.4s, {w1:v}.16b, {x1:v}.16b",
6202                    a0 = inout(vreg) a0, a1 = inout(vreg) a1,
6203                    w0 = in(vreg) $w0, x0 = in(vreg) x0, w1 = in(vreg) $w1, x1 = in(vreg) x1,
6204                    options(pure, nomem, nostack),
6205                );
6206                vaddq_s32(a0, a1)
6207            }};
6208        }
6209        let base = bytes.as_ptr().add(r * gpr * Q1_TILE);
6210        let row_base = r * gpr * Q1_TILE;
6211        let abs_end = bytes.len();
6212        let mut accv = [vdupq_n_f32(0.0); 4];
6213        let mut gi = 0;
6214        while gi + 4 <= gpr && row_base + (gi + 4) * Q1_TILE + 4 <= abs_end {
6215            let t0 = base.add(gi * Q1_TILE);
6216            let ld_a = vld1q_u8(t0);
6217            let ld_b = vld1q_u8(t0.add(2 * Q1_TILE));
6218            // Unpack ONCE — eight ±mask vectors serve all four streams.
6219            let w00 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8(ld_a, iw00), m));
6220            let w01 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8(ld_a, iw01), m));
6221            let w10 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8(ld_a, iw10), m));
6222            let w11 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8(ld_a, iw11), m));
6223            let w20 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8(ld_b, iw00), m));
6224            let w21 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8(ld_b, iw01), m));
6225            let w30 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8(ld_b, iw10), m));
6226            let w31 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8(ld_b, iw11), m));
6227            let sc16 = vqtbl2_u8(uint8x16x2_t(ld_a, ld_b), isc);
6228            let scf: float32x4_t;
6229            asm!(
6230                "fcvtl {o:v}.4s, {i:v}.4h",
6231                o = out(vreg) scf, i = in(vreg) sc16,
6232                options(pure, nomem, nostack),
6233            );
6234            for k in 0..4 {
6235                let xp = xs[k].as_ptr();
6236                let d0 = sdot2!(w00, w01, xp.add(gi * GROUP_SIZE));
6237                let d1 = sdot2!(w10, w11, xp.add((gi + 1) * GROUP_SIZE));
6238                let d2 = sdot2!(w20, w21, xp.add((gi + 2) * GROUP_SIZE));
6239                let d3 = sdot2!(w30, w31, xp.add((gi + 3) * GROUP_SIZE));
6240                let neg = vpaddq_s32(vpaddq_s32(d0, d1), vpaddq_s32(d2, d3));
6241                let g = vld1q_s32(gs[k].as_ptr().add(gi));
6242                let dots = vnegq_s32(vaddq_s32(vshlq_n_s32::<1>(neg), g));
6243                accv[k] = vfmaq_f32(accv[k], vcvtq_f32_s32(dots), scf);
6244            }
6245            gi += 4;
6246        }
6247        let mut acc = [
6248            vaddvq_f32(accv[0]),
6249            vaddvq_f32(accv[1]),
6250            vaddvq_f32(accv[2]),
6251            vaddvq_f32(accv[3]),
6252        ];
6253        while gi < gpr {
6254            let t = base.add(gi * Q1_TILE);
6255            let sc = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
6256            let v0 = vcombine_u8(vdup_n_u8(*t.add(2)), vdup_n_u8(*t.add(3)));
6257            let v1 = vcombine_u8(vdup_n_u8(*t.add(4)), vdup_n_u8(*t.add(5)));
6258            let w0 = vreinterpretq_s8_u8(vtstq_u8(v0, m));
6259            let w1 = vreinterpretq_s8_u8(vtstq_u8(v1, m));
6260            for k in 0..4 {
6261                let d = vaddvq_s32(sdot2!(w0, w1, xs[k].as_ptr().add(gi * GROUP_SIZE)));
6262                acc[k] += (-(2 * d + *gs[k].as_ptr().add(gi))) as f32 * sc;
6263            }
6264            gi += 1;
6265        }
6266        acc
6267    }
6268}
6269
6270/// (weight ±1, scale) of one q1 element — the exact outlier term.
6271#[inline]
6272fn q1_outlier(bytes: &[u8], r: usize, gpr: usize, j: usize) -> (f32, f32) {
6273    let gi = j / GROUP_SIZE;
6274    let k = j % GROUP_SIZE;
6275    let tile = &bytes[(r * gpr + gi) * Q1_TILE..(r * gpr + gi + 1) * Q1_TILE];
6276    let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
6277    let bit = (tile[2 + k / 8] >> (k % 8)) & 1;
6278    ((bit as i32 * 2 - 1) as f32, s)
6279}
6280
6281/// Exact scalar q1 row (CMF_SDOT=0 contract).
6282#[inline]
6283fn q1_row_exact(bytes: &[u8], r: usize, gpr: usize, x: &[f32]) -> f32 {
6284    let mut acc = 0f32;
6285    for gi in 0..gpr {
6286        let tile = &bytes[(r * gpr + gi) * Q1_TILE..(r * gpr + gi + 1) * Q1_TILE];
6287        let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
6288        let xg = &x[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
6289        let mut ga = 0f32;
6290        for (j, &b) in tile[2..].iter().enumerate() {
6291            for k in 0..8 {
6292                ga += (((b >> k) & 1) as f32 * 2.0 - 1.0) * xg[j * 8 + k];
6293            }
6294        }
6295        acc += ga * s;
6296    }
6297    acc
6298}
6299
6300/// One q1 row range via A8W8 (the body of `q1_matvec`'s hot loop,
6301/// extracted so multi-matrix jobs drive the same kernel).
6302#[allow(clippy::too_many_arguments)]
6303fn q1_range_a8w8(
6304    bytes: &[u8],
6305    gpr: usize,
6306    act: &SplitAct,
6307    gsum: &[i32],
6308    out: SendMut,
6309    start: usize,
6310    end: usize,
6311) {
6312    for r in start..end {
6313        let mut acc = dot_q1_row_i8(bytes, r, gpr, &act.xq, gsum) * act.sx;
6314        for &(j, xv) in &act.outliers {
6315            let (w, s) = q1_outlier(bytes, r, gpr, j);
6316            acc += w * s * xv;
6317        }
6318        // SAFETY: disjoint row ranges per worker.
6319        unsafe { *out.at(r) = acc };
6320    }
6321}
6322
6323/// Exact-scalar q1 row range (CMF_SDOT=0 contract).
6324fn q1_range_f32(bytes: &[u8], gpr: usize, x: &[f32], out: SendMut, start: usize, end: usize) {
6325    for r in start..end {
6326        // SAFETY: disjoint row ranges per worker.
6327        unsafe { *out.at(r) = q1_row_exact(bytes, r, gpr, x) };
6328    }
6329}
6330
6331/// q1t per-row overlay locator. After the base (`base_len`) come
6332/// `[u32 row_ptr[rows+1]]` then `[(u16 col, f16 val)]` grouped by row (row
6333/// `r`'s entries are `[row_ptr[r], row_ptr[r+1])`). Returns
6334/// `(row_ptr offset, entries offset, present)`.
6335fn q1t_overlay(bytes: &[u8], base_len: usize, rows: usize) -> (usize, usize, bool) {
6336    let entries = base_len + (rows + 1) * 4;
6337    (base_len, entries, entries <= bytes.len())
6338}
6339
6340/// Read `row_ptr[r]` from the overlay's prefix-sum table.
6341#[inline]
6342fn q1t_rowptr(bytes: &[u8], rp_off: usize, r: usize) -> usize {
6343    let o = rp_off + r * 4;
6344    u32::from_le_bytes([bytes[o], bytes[o + 1], bytes[o + 2], bytes[o + 3]]) as usize
6345}
6346
6347/// Byte → the 5 ternary signs it packs `{−1,0,+1}` as f32, precomputed so
6348/// decoding a q1t code is a table load, not the base-3 divide/modulo per
6349/// weight (division is ~20–40× the cost of a load). Built at compile time.
6350const SIGN5: [[f32; 5]; 256] = {
6351    let mut lut = [[0.0f32; 5]; 256];
6352    let pow3 = [1u16, 3, 9, 27, 81];
6353    let mut byte = 0usize;
6354    while byte < 256 {
6355        let mut i = 0usize;
6356        while i < 5 {
6357            let code = (byte as u16 / pow3[i]) % 3;
6358            lut[byte][i] = if code == 1 {
6359                1.0
6360            } else if code == 2 {
6361                -1.0
6362            } else {
6363                0.0
6364            };
6365            i += 1;
6366        }
6367        byte += 1;
6368    }
6369    lut
6370};
6371
6372/// Same table, as i8 signs — the operand for the int8 SDOT base kernel.
6373const SIGN5_I8: [[i8; 5]; 256] = {
6374    let mut lut = [[0i8; 5]; 256];
6375    let pow3 = [1u16, 3, 9, 27, 81];
6376    let mut byte = 0usize;
6377    while byte < 256 {
6378        let mut i = 0usize;
6379        while i < 5 {
6380            let code = (byte as u16 / pow3[i]) % 3;
6381            lut[byte][i] = if code == 1 {
6382                1
6383            } else if code == 2 {
6384                -1
6385            } else {
6386                0
6387            };
6388            i += 1;
6389        }
6390        byte += 1;
6391    }
6392    lut
6393};
6394
6395/// The same 5 i8 signs packed into a u64 (`[s0 s1 s2 s3 s4 0 0 0]`, LE) so the
6396/// group unpack is 7 unaligned u64 stores at offsets 0,5,10,…,30 instead of
6397/// six 5-byte copies + LUT indexing — each store's trailing zeros are fixed by
6398/// the next store, and the last one runs 6 B past the 32nd weight (the unpack
6399/// buffer is padded to 40). This is the decode/prefill hot inner op.
6400const SIGN5_U64: [u64; 256] = {
6401    let mut lut = [0u64; 256];
6402    let pow3 = [1u16, 3, 9, 27, 81];
6403    let mut byte = 0usize;
6404    while byte < 256 {
6405        let mut v = 0u64;
6406        let mut i = 0usize;
6407        while i < 5 {
6408            let code = (byte as u16 / pow3[i]) % 3;
6409            let s: u8 = if code == 1 {
6410                1
6411            } else if code == 2 {
6412                0xFF
6413            } else {
6414                0
6415            };
6416            v |= (s as u64) << (i * 8);
6417            i += 1;
6418        }
6419        lut[byte] = v;
6420        byte += 1;
6421    }
6422    lut
6423};
6424
6425/// Ternary base weight at `(row r, col j)` = `sign(code)·s_group`. Used to add
6426/// back activation-outlier columns, whose `x` was zeroed for the int8 bulk dot
6427/// (`split_act`). At a weight-outlier position the code is 0, so this is 0 and
6428/// the overlay correction owns that column — no double counting.
6429#[inline]
6430fn q1t_base_weight(bytes: &[u8], r: usize, gpr: usize, j: usize) -> f32 {
6431    const TILE: usize = cortiq_core::quant::Q1T_TILE;
6432    let off = (r * gpr + j / GROUP_SIZE) * TILE;
6433    let s = f16_to_f32(u16::from_le_bytes([bytes[off], bytes[off + 1]]));
6434    let within = j % GROUP_SIZE;
6435    SIGN5[bytes[off + 2 + within / 5] as usize][within % 5] * s
6436}
6437
6438/// One 32-group int8 dot via two SDOTs. Bit-exact vs the scalar i8 sum
6439/// (integer accumulation is order-independent).
6440#[cfg(target_arch = "aarch64")]
6441#[target_feature(enable = "neon,dotprod")]
6442#[inline]
6443unsafe fn sdot32_i8(w: *const i8, x: *const i8) -> i32 {
6444    // SAFETY: caller guarantees 32 readable i8 at each pointer.
6445    unsafe {
6446        use core::arch::aarch64::*;
6447        use core::arch::asm;
6448        let w0 = vld1q_s8(w);
6449        let w1 = vld1q_s8(w.add(16));
6450        let x0 = vld1q_s8(x);
6451        let x1 = vld1q_s8(x.add(16));
6452        let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
6453        asm!(
6454            "sdot {a0:v}.4s, {w0:v}.16b, {x0:v}.16b",
6455            "sdot {a1:v}.4s, {w1:v}.16b, {x1:v}.16b",
6456            a0 = inout(vreg) a0, a1 = inout(vreg) a1,
6457            w0 = in(vreg) w0, x0 = in(vreg) x0, w1 = in(vreg) w1, x1 = in(vreg) x1,
6458            options(pure, nomem, nostack),
6459        );
6460        vaddvq_s32(vaddq_s32(a0, a1))
6461    }
6462}
6463
6464/// One 32-group int8 dot via AVX2: signed·signed as `maddubs(|w|, sign(x,w))`
6465/// then `madd` and a horizontal reduce (the same idiom as `dot_q4t_row_avx2`).
6466#[cfg(target_arch = "x86_64")]
6467#[target_feature(enable = "avx2")]
6468#[inline]
6469unsafe fn i8dot32_avx2(w: *const i8, x: *const i8) -> i32 {
6470    // SAFETY: caller guarantees 32 readable i8 at each pointer.
6471    unsafe {
6472        use core::arch::x86_64::*;
6473        let wv = _mm256_loadu_si256(w as *const __m256i);
6474        let xv = _mm256_loadu_si256(x as *const __m256i);
6475        let p16 = _mm256_maddubs_epi16(_mm256_abs_epi8(wv), _mm256_sign_epi8(xv, wv));
6476        let d = _mm256_madd_epi16(p16, _mm256_set1_epi16(1));
6477        let hi128 = _mm256_extracti128_si256::<1>(d);
6478        let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
6479        let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
6480        let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
6481        _mm_cvtsi128_si32(s32)
6482    }
6483}
6484
6485/// Unpack one q1t group's base-3 codes into 32 i8 signs via 7 unaligned u64
6486/// stores (see `SIGN5_U64`). `dst` MUST have ≥ 40 bytes: the 7th store writes
6487/// `dst[30..38]`. Stores go in order so each one's trailing zeros are
6488/// overwritten by the next; the final 6 padding bytes are unused by the dot.
6489#[inline]
6490fn q1t_unpack_group_i8(codes: *const u8, dst: &mut [i8]) {
6491    debug_assert!(dst.len() >= 40);
6492    // SAFETY: codes points at 7 readable bytes; dst has ≥ 40 bytes so every
6493    // 8-byte store at offset bi*5 (bi ≤ 6 → ≤ 30) stays in bounds.
6494    unsafe {
6495        let p = dst.as_mut_ptr();
6496        for bi in 0..7 {
6497            core::ptr::write_unaligned(
6498                p.add(bi * 5) as *mut u64,
6499                SIGN5_U64[*codes.add(bi) as usize],
6500            );
6501        }
6502    }
6503}
6504
6505/// One 32-group int8 dot, arch-dispatched (the matmat inner loop, where the
6506/// row's signs are unpacked once and dotted against every batch input).
6507/// Callers are gated by `a8w8_enabled()`, so the target-feature arms are
6508/// reachable; the scalar arm is a non-SIMD-arch fallback.
6509#[inline]
6510fn q1t_i8dot32(w: *const i8, x: *const i8) -> i32 {
6511    #[cfg(target_arch = "aarch64")]
6512    unsafe {
6513        return sdot32_i8(w, x);
6514    }
6515    #[cfg(target_arch = "x86_64")]
6516    unsafe {
6517        return i8dot32_avx2(w, x);
6518    }
6519    #[allow(unreachable_code)]
6520    unsafe {
6521        let mut s = 0i32;
6522        for k in 0..GROUP_SIZE {
6523            s += *w.add(k) as i32 * *x.add(k) as i32;
6524        }
6525        s
6526    }
6527}
6528
6529#[inline]
6530unsafe fn q1t_unpack_reg_u64s(codes: *const u8) -> (u64, u64, u64, u64) {
6531    let (s0, s1, s2, s3, s4, s5, s6) = unsafe {
6532        (
6533            SIGN5_U64[*codes as usize],
6534            SIGN5_U64[*codes.add(1) as usize],
6535            SIGN5_U64[*codes.add(2) as usize],
6536            SIGN5_U64[*codes.add(3) as usize],
6537            SIGN5_U64[*codes.add(4) as usize],
6538            SIGN5_U64[*codes.add(5) as usize],
6539            SIGN5_U64[*codes.add(6) as usize],
6540        )
6541    };
6542
6543    let u0 = s0 | (s1 << 40);
6544    let u1 = (s1 >> 24) | (s2 << 16) | (s3 << 56);
6545    let u2 = (s3 >> 8) | (s4 << 32);
6546    let u3 = (s4 >> 32) | (s5 << 8) | (s6 << 48);
6547
6548    (u0, u1, u2, u3)
6549}
6550
6551/// One q1t row's int8 base dot: `Σ_group s·dot(signs, xq)` (before the shared
6552/// `sx`). Direct register unpacking (zero stack stores/loads, no STLF stalls).
6553/// ARM SDOT.
6554#[cfg(target_arch = "aarch64")]
6555#[target_feature(enable = "neon,dotprod")]
6556unsafe fn q1t_dot_row_sdot(bytes: &[u8], r: usize, gpr: usize, xq: &[i8]) -> f32 {
6557    use core::arch::aarch64::*;
6558    use core::arch::asm;
6559    unsafe {
6560        const TILE: usize = cortiq_core::quant::Q1T_TILE;
6561        let mut acc = 0f32;
6562        let bytes_ptr = bytes.as_ptr();
6563        let xq_ptr = xq.as_ptr();
6564        let row_off = r * gpr * TILE;
6565
6566        let gpr2 = gpr & !1;
6567        let mut gi = 0;
6568        while gi < gpr2 {
6569            let off0 = row_off + gi * TILE;
6570            let off1 = off0 + TILE;
6571            let s0 = f16_to_f32(u16::from_le_bytes([
6572                *bytes_ptr.add(off0),
6573                *bytes_ptr.add(off0 + 1),
6574            ]));
6575            let s1 = f16_to_f32(u16::from_le_bytes([
6576                *bytes_ptr.add(off1),
6577                *bytes_ptr.add(off1 + 1),
6578            ]));
6579
6580            let (u0_0, u1_0, u2_0, u3_0) = q1t_unpack_reg_u64s(bytes_ptr.add(off0 + 2));
6581            let (u0_1, u1_1, u2_1, u3_1) = q1t_unpack_reg_u64s(bytes_ptr.add(off1 + 2));
6582
6583            let w0_0 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u0_0), vcreate_u64(u1_0)));
6584            let w1_0 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u2_0), vcreate_u64(u3_0)));
6585            let w0_1 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u0_1), vcreate_u64(u1_1)));
6586            let w1_1 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u2_1), vcreate_u64(u3_1)));
6587
6588            let x0_0 = vld1q_s8(xq_ptr.add(gi * GROUP_SIZE));
6589            let x1_0 = vld1q_s8(xq_ptr.add(gi * GROUP_SIZE + 16));
6590            let x0_1 = vld1q_s8(xq_ptr.add((gi + 1) * GROUP_SIZE));
6591            let x1_1 = vld1q_s8(xq_ptr.add((gi + 1) * GROUP_SIZE + 16));
6592
6593            let (mut a0_0, mut a1_0) = (vdupq_n_s32(0), vdupq_n_s32(0));
6594            let (mut a0_1, mut a1_1) = (vdupq_n_s32(0), vdupq_n_s32(0));
6595            asm!(
6596                "sdot {a0_0:v}.4s, {w0_0:v}.16b, {x0_0:v}.16b",
6597                "sdot {a1_0:v}.4s, {w1_0:v}.16b, {x1_0:v}.16b",
6598                "sdot {a0_1:v}.4s, {w0_1:v}.16b, {x0_1:v}.16b",
6599                "sdot {a1_1:v}.4s, {w1_1:v}.16b, {x1_1:v}.16b",
6600                a0_0 = inout(vreg) a0_0, a1_0 = inout(vreg) a1_0,
6601                a0_1 = inout(vreg) a0_1, a1_1 = inout(vreg) a1_1,
6602                w0_0 = in(vreg) w0_0, x0_0 = in(vreg) x0_0, w1_0 = in(vreg) w1_0, x1_0 = in(vreg) x1_0,
6603                w0_1 = in(vreg) w0_1, x0_1 = in(vreg) x0_1, w1_1 = in(vreg) w1_1, x1_1 = in(vreg) x1_1,
6604                options(pure, nomem, nostack),
6605            );
6606            let d0 = vaddvq_s32(vaddq_s32(a0_0, a1_0));
6607            let d1 = vaddvq_s32(vaddq_s32(a0_1, a1_1));
6608            acc += d0 as f32 * s0 + d1 as f32 * s1;
6609            gi += 2;
6610        }
6611
6612        if gi < gpr {
6613            let off = row_off + gi * TILE;
6614            let s = f16_to_f32(u16::from_le_bytes([
6615                *bytes_ptr.add(off),
6616                *bytes_ptr.add(off + 1),
6617            ]));
6618            let (u0, u1, u2, u3) = q1t_unpack_reg_u64s(bytes_ptr.add(off + 2));
6619            let w0 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u0), vcreate_u64(u1)));
6620            let w1 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u2), vcreate_u64(u3)));
6621            let x0 = vld1q_s8(xq_ptr.add(gi * GROUP_SIZE));
6622            let x1 = vld1q_s8(xq_ptr.add(gi * GROUP_SIZE + 16));
6623            let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
6624            asm!(
6625                "sdot {a0:v}.4s, {w0:v}.16b, {x0:v}.16b",
6626                "sdot {a1:v}.4s, {w1:v}.16b, {x1:v}.16b",
6627                a0 = inout(vreg) a0, a1 = inout(vreg) a1,
6628                w0 = in(vreg) w0, x0 = in(vreg) x0, w1 = in(vreg) w1, x1 = in(vreg) x1,
6629                options(pure, nomem, nostack),
6630            );
6631            let d = vaddvq_s32(vaddq_s32(a0, a1));
6632            acc += d as f32 * s;
6633        }
6634        acc
6635    }
6636}
6637
6638/// x86 AVX2 mirror of `q1t_dot_row_sdot` (maddubs int8 dot per group).
6639#[cfg(target_arch = "x86_64")]
6640#[target_feature(enable = "avx2")]
6641unsafe fn q1t_dot_row_avx2(bytes: &[u8], r: usize, gpr: usize, xq: &[i8]) -> f32 {
6642    use core::arch::x86_64::*;
6643    unsafe {
6644        const TILE: usize = cortiq_core::quant::Q1T_TILE;
6645        let mut acc = 0f32;
6646        let bytes_ptr = bytes.as_ptr();
6647        let xq_ptr = xq.as_ptr();
6648        let row_off = r * gpr * TILE;
6649
6650        let ones = _mm256_set1_epi16(1);
6651        for gi in 0..gpr {
6652            let off = row_off + gi * TILE;
6653            let s = f16_to_f32(u16::from_le_bytes([
6654                *bytes_ptr.add(off),
6655                *bytes_ptr.add(off + 1),
6656            ]));
6657            let (u0, u1, u2, u3) = q1t_unpack_reg_u64s(bytes_ptr.add(off + 2));
6658            let wv = _mm256_set_epi64x(u3 as i64, u2 as i64, u1 as i64, u0 as i64);
6659            let xv = _mm256_loadu_si256(xq_ptr.add(gi * GROUP_SIZE) as *const __m256i);
6660            let p16 = _mm256_maddubs_epi16(_mm256_abs_epi8(wv), _mm256_sign_epi8(xv, wv));
6661            let d256 = _mm256_madd_epi16(p16, ones);
6662            let d128 = _mm_add_epi32(
6663                _mm256_castsi256_si128(d256),
6664                _mm256_extracti128_si256(d256, 1),
6665            );
6666            let d64 = _mm_add_epi32(d128, _mm_shuffle_epi32(d128, 0xee));
6667            let d32 = _mm_cvtsi128_si32(_mm_add_epi32(d64, _mm_shuffle_epi32(d64, 0x55)));
6668            acc += d32 as f32 * s;
6669        }
6670        acc
6671    }
6672}
6673
6674/// VNNI twin of `q1t_dot_row_avx2` (see `dpbusd_hsum`).
6675#[cfg(target_arch = "x86_64")]
6676#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
6677unsafe fn q1t_dot_row_vnni(bytes: &[u8], r: usize, gpr: usize, xq: &[i8]) -> f32 {
6678    use core::arch::x86_64::*;
6679    // SAFETY: same tile/xq contracts as `q1t_dot_row_avx2`.
6680    unsafe {
6681        const TILE: usize = cortiq_core::quant::Q1T_TILE;
6682        let mut acc = 0f32;
6683        let bytes_ptr = bytes.as_ptr();
6684        let xq_ptr = xq.as_ptr();
6685        let row_off = r * gpr * TILE;
6686        for gi in 0..gpr {
6687            let off = row_off + gi * TILE;
6688            let s = f16_to_f32(u16::from_le_bytes([
6689                *bytes_ptr.add(off),
6690                *bytes_ptr.add(off + 1),
6691            ]));
6692            let (u0, u1, u2, u3) = q1t_unpack_reg_u64s(bytes_ptr.add(off + 2));
6693            let wv = _mm256_set_epi64x(u3 as i64, u2 as i64, u1 as i64, u0 as i64);
6694            let xv = _mm256_loadu_si256(xq_ptr.add(gi * GROUP_SIZE) as *const __m256i);
6695            let d = dpbusd_hsum(_mm256_abs_epi8(wv), _mm256_sign_epi8(xv, wv));
6696            acc += d as f32 * s;
6697        }
6698        acc
6699    }
6700}
6701
6702/// Per-row int8 base dot, dispatched once per row (matvec decode hot path).
6703/// Callers are gated by `a8w8_enabled()`, so the target-feature kernels are
6704/// reachable.
6705#[inline]
6706fn q1t_dot_row_i8(bytes: &[u8], r: usize, gpr: usize, xq: &[i8]) -> f32 {
6707    #[cfg(target_arch = "aarch64")]
6708    unsafe {
6709        return q1t_dot_row_sdot(bytes, r, gpr, xq);
6710    }
6711    #[cfg(target_arch = "x86_64")]
6712    unsafe {
6713        if vnni_tiles_enabled() {
6714            return q1t_dot_row_vnni(bytes, r, gpr, xq);
6715        }
6716        return q1t_dot_row_avx2(bytes, r, gpr, xq);
6717    }
6718    #[allow(unreachable_code)]
6719    {
6720        const TILE: usize = cortiq_core::quant::Q1T_TILE;
6721        let mut acc = 0f32;
6722        let mut sg = [0i8; GROUP_SIZE + 8]; // +8 slack for the u64-store unpack
6723        for gi in 0..gpr {
6724            let off = (r * gpr + gi) * TILE;
6725            let s = f16_to_f32(u16::from_le_bytes([bytes[off], bytes[off + 1]]));
6726            q1t_unpack_group_i8(bytes.as_ptr().wrapping_add(off + 2), &mut sg);
6727            let mut d = 0i32;
6728            for k in 0..GROUP_SIZE {
6729                d += sg[k] as i32 * xq[gi * GROUP_SIZE + k] as i32;
6730            }
6731            acc += d as f32 * s;
6732        }
6733        acc
6734    }
6735}
6736
6737/// Σ over a row's outliers of `value·x[col]` — the correction that adds the
6738/// overlay's exact weights on top of the base dot. INVARIANT: the encoder
6739/// writes ternary code 0 at every outlier position (`quantize_q1t`), so the
6740/// base contributes nothing there and this is a plain `value·x`, not
6741/// `(value − base)·x` — no scattered per-outlier scale read. Row `r`'s entries
6742/// are the contiguous slice `[row_ptr[r], row_ptr[r+1])`, so no binary search.
6743fn q1t_row_outlier_correction(
6744    bytes: &[u8],
6745    r: usize,
6746    rp_off: usize,
6747    entries_off: usize,
6748    has_ov: bool,
6749    x: &[f32],
6750) -> f32 {
6751    if !has_ov {
6752        return 0.0;
6753    }
6754    let (c0, c1) = (
6755        q1t_rowptr(bytes, rp_off, r),
6756        q1t_rowptr(bytes, rp_off, r + 1),
6757    );
6758    let mut corr = 0f32;
6759    for p in c0..c1 {
6760        let e = entries_off + p * 4;
6761        let col = u16::from_le_bytes([bytes[e], bytes[e + 1]]) as usize;
6762        let val = f16_to_f32(u16::from_le_bytes([bytes[e + 2], bytes[e + 3]]));
6763        corr += val * x[col];
6764    }
6765    corr
6766}
6767
6768/// Dequantize one q1t row into `buf[..cols]` via the sign LUT (no division),
6769/// then apply the row's outliers (its `[row_ptr[r], row_ptr[r+1])` slice).
6770/// Used by the batched (prefill) path where the decode amortizes over the batch.
6771fn q1t_dequant_row(
6772    bytes: &[u8],
6773    r: usize,
6774    gpr: usize,
6775    rp_off: usize,
6776    entries_off: usize,
6777    has_ov: bool,
6778    buf: &mut [f32],
6779) {
6780    const TILE: usize = cortiq_core::quant::Q1T_TILE;
6781    for g in 0..gpr {
6782        let off = (r * gpr + g) * TILE;
6783        let s = f16_to_f32(u16::from_le_bytes([bytes[off], bytes[off + 1]]));
6784        let codes = &bytes[off + 2..off + TILE];
6785        let bc = g * GROUP_SIZE;
6786        // 6 full bytes (30 codes) + a 7th byte holding the last 2.
6787        for bi in 0..6 {
6788            let lut = &SIGN5[codes[bi] as usize];
6789            let d = &mut buf[bc + bi * 5..bc + bi * 5 + 5];
6790            for i in 0..5 {
6791                d[i] = lut[i] * s;
6792            }
6793        }
6794        let lut = &SIGN5[codes[6] as usize];
6795        buf[bc + 30] = lut[0] * s;
6796        buf[bc + 31] = lut[1] * s;
6797    }
6798    if !has_ov {
6799        return;
6800    }
6801    let (c0, c1) = (
6802        q1t_rowptr(bytes, rp_off, r),
6803        q1t_rowptr(bytes, rp_off, r + 1),
6804    );
6805    for p in c0..c1 {
6806        let e = entries_off + p * 4;
6807        let col = u16::from_le_bytes([bytes[e], bytes[e + 1]]) as usize;
6808        buf[col] = f16_to_f32(u16::from_le_bytes([bytes[e + 2], bytes[e + 3]]));
6809    }
6810}
6811
6812/// Add the sparse outlier overlay onto a base dot already in `out` (the GPU
6813/// computes the ternary base; the overlay stays on the CPU — its entries are
6814/// few and its per-row gather doesn't vectorize on the GPU). Row-parallel.
6815fn q1t_add_overlay(
6816    bytes: &[u8],
6817    x: &[f32],
6818    rows: usize,
6819    cols: usize,
6820    out: &mut [f32],
6821    pool: Option<&Pool>,
6822) {
6823    const TILE: usize = cortiq_core::quant::Q1T_TILE;
6824    let gpr = cols / GROUP_SIZE;
6825    let (rp_off, ent_off, has_ov) = q1t_overlay(bytes, rows * gpr * TILE, rows);
6826    if !has_ov {
6827        return;
6828    }
6829    let out_addr = SendMut(out.as_mut_ptr());
6830    let run = move |start: usize, end: usize| {
6831        for r in start..end {
6832            let corr = q1t_row_outlier_correction(bytes, r, rp_off, ent_off, has_ov, x);
6833            // SAFETY: disjoint rows; add onto the base the GPU already wrote.
6834            unsafe { *out_addr.at(r) += corr };
6835        }
6836    };
6837    dispatch_rows(pool, rows, &run);
6838}
6839
6840/// Q1T row range via the A8W8 int8 path — shared activation split,
6841/// per-row: base SDOT dot + outlier correction + overlay.
6842#[allow(clippy::too_many_arguments)]
6843fn q1t_range_a8w8(
6844    bytes: &[u8],
6845    gpr: usize,
6846    rp_off: usize,
6847    ent_off: usize,
6848    has_ov: bool,
6849    act: &SplitAct,
6850    x: &[f32],
6851    out: SendMut,
6852    start: usize,
6853    end: usize,
6854) {
6855    for r in start..end {
6856        let mut acc = q1t_dot_row_i8(bytes, r, gpr, &act.xq) * act.sx;
6857        for &(j, xv) in &act.outliers {
6858            acc += q1t_base_weight(bytes, r, gpr, j) * xv;
6859        }
6860        acc += q1t_row_outlier_correction(bytes, r, rp_off, ent_off, has_ov, x);
6861        // SAFETY: disjoint row ranges per worker.
6862        unsafe { *out.at(r) = acc };
6863    }
6864}
6865
6866/// Q1T row range via the f32 path (no SDOT) — for matvec_many batched
6867/// dispatch when a8w8 is unavailable.
6868#[allow(clippy::too_many_arguments)]
6869fn q1t_range_f32_batch(
6870    bytes: &[u8],
6871    gpr: usize,
6872    rp_off: usize,
6873    ent_off: usize,
6874    has_ov: bool,
6875    x: &[f32],
6876    out: SendMut,
6877    start: usize,
6878    end: usize,
6879) {
6880    const TILE: usize = cortiq_core::quant::Q1T_TILE;
6881    let mut sg = [0f32; GROUP_SIZE];
6882    for r in start..end {
6883        let mut acc = 0f32;
6884        for g in 0..gpr {
6885            let off = (r * gpr + g) * TILE;
6886            let s = f16_to_f32(u16::from_le_bytes([bytes[off], bytes[off + 1]]));
6887            let codes = &bytes[off + 2..off + TILE];
6888            let xg = &x[g * GROUP_SIZE..g * GROUP_SIZE + GROUP_SIZE];
6889            for bi in 0..6 {
6890                sg[bi * 5..bi * 5 + 5].copy_from_slice(&SIGN5[codes[bi] as usize]);
6891            }
6892            let lut = &SIGN5[codes[6] as usize];
6893            sg[30] = lut[0];
6894            sg[31] = lut[1];
6895            let mut gsum = 0f32;
6896            for k in 0..GROUP_SIZE {
6897                gsum += sg[k] * xg[k];
6898            }
6899            acc += s * gsum;
6900        }
6901        acc += q1t_row_outlier_correction(bytes, r, rp_off, ent_off, has_ov, x);
6902        // SAFETY: disjoint row ranges per worker.
6903        unsafe { *out.at(r) = acc };
6904    }
6905}
6906
6907/// Ternary (q1t) matvec — decode+dot straight from mmap, one group at a time:
6908/// no per-ROW buffer, no division (the sign LUT), and a tiny per-group sign
6909/// buffer so the 32-wide dot vectorizes. This is the decode hot path.
6910fn q1t_matvec(
6911    bytes: &[u8],
6912    x: &[f32],
6913    rows: usize,
6914    cols: usize,
6915    out: &mut [f32],
6916    pool: Option<&Pool>,
6917) {
6918    debug_assert_eq!(out.len(), rows);
6919    const TILE: usize = cortiq_core::quant::Q1T_TILE;
6920    let gpr = cols / GROUP_SIZE;
6921    let (rp_off, ent_off, has_ov) = q1t_overlay(bytes, rows * gpr * TILE, rows);
6922    let out_addr = SendMut(out.as_mut_ptr());
6923    // int8 SDOT base dot (ARM dotprod): ~4× the f32 arithmetic. x → i8 once
6924    // (`split_act`), activation outliers added back exactly in f32, weight
6925    // overlay on top. ARM SDOT / x86 AVX2; CMF_SDOT=0 keeps the exact f32 path.
6926    if a8w8_enabled() {
6927        let act = split_act(x);
6928        let act = &act;
6929        let run = move |start: usize, end: usize| {
6930            for r in start..end {
6931                let mut acc = q1t_dot_row_i8(bytes, r, gpr, &act.xq) * act.sx;
6932                for &(j, xv) in &act.outliers {
6933                    acc += q1t_base_weight(bytes, r, gpr, j) * xv;
6934                }
6935                acc += q1t_row_outlier_correction(bytes, r, rp_off, ent_off, has_ov, x);
6936                // SAFETY: disjoint row ranges per worker.
6937                unsafe { *out_addr.at(r) = acc };
6938            }
6939        };
6940        dispatch_rows(pool, rows, &run);
6941        return;
6942    }
6943    let run = move |start: usize, end: usize| {
6944        // Per-group signs, unpacked contiguously so the dot below is a clean
6945        // 32-wide reduction the autovectorizer turns into f32x4 FMAs — the
6946        // 5-values-per-byte base-3 layout won't SIMD in place.
6947        let mut sg = [0f32; GROUP_SIZE];
6948        for r in start..end {
6949            let mut acc = 0f32;
6950            for g in 0..gpr {
6951                let off = (r * gpr + g) * TILE;
6952                let s = f16_to_f32(u16::from_le_bytes([bytes[off], bytes[off + 1]]));
6953                let codes = &bytes[off + 2..off + TILE];
6954                let xg = &x[g * GROUP_SIZE..g * GROUP_SIZE + GROUP_SIZE];
6955                for bi in 0..6 {
6956                    sg[bi * 5..bi * 5 + 5].copy_from_slice(&SIGN5[codes[bi] as usize]);
6957                }
6958                let lut = &SIGN5[codes[6] as usize];
6959                sg[30] = lut[0];
6960                sg[31] = lut[1];
6961                let mut gsum = 0f32;
6962                for k in 0..GROUP_SIZE {
6963                    gsum += sg[k] * xg[k];
6964                }
6965                acc += s * gsum;
6966            }
6967            acc += q1t_row_outlier_correction(bytes, r, rp_off, ent_off, has_ov, x);
6968            unsafe { *out_addr.at(r) = acc };
6969        }
6970    };
6971    dispatch_rows(pool, rows, &run);
6972}
6973
6974/// Fused-pair twin of `q1t_dot_row_sdot`: ONE register unpack of the
6975/// ternary codes serves BOTH activation streams (the unpack chain is
6976/// the dominant per-row cost — MTP verify pairs paid it twice). Per
6977/// stream the group order and f32 accumulation match the single-row
6978/// kernel exactly, so pair == 2×matvec bit-for-bit.
6979#[cfg(target_arch = "aarch64")]
6980#[target_feature(enable = "neon,dotprod")]
6981unsafe fn q1t_dot_row_sdot2(bytes: &[u8], r: usize, gpr: usize, xa: &[i8], xb: &[i8]) -> [f32; 2] {
6982    use core::arch::aarch64::*;
6983    use core::arch::asm;
6984    // SAFETY: same slice-length contracts as `q1t_dot_row_sdot`, ×2.
6985    unsafe {
6986        const TILE: usize = cortiq_core::quant::Q1T_TILE;
6987        let bytes_ptr = bytes.as_ptr();
6988        let row_off = r * gpr * TILE;
6989        let xp = [xa.as_ptr(), xb.as_ptr()];
6990        let mut acc = [0f32; 2];
6991        macro_rules! sdot2 {
6992            ($w0:expr, $w1:expr, $x:expr) => {{
6993                let x0 = vld1q_s8($x);
6994                let x1 = vld1q_s8($x.add(16));
6995                let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
6996                asm!(
6997                    "sdot {a0:v}.4s, {w0:v}.16b, {x0:v}.16b",
6998                    "sdot {a1:v}.4s, {w1:v}.16b, {x1:v}.16b",
6999                    a0 = inout(vreg) a0, a1 = inout(vreg) a1,
7000                    w0 = in(vreg) $w0, x0 = in(vreg) x0, w1 = in(vreg) $w1, x1 = in(vreg) x1,
7001                    options(pure, nomem, nostack),
7002                );
7003                vaddvq_s32(vaddq_s32(a0, a1))
7004            }};
7005        }
7006        let gpr2 = gpr & !1;
7007        let mut gi = 0;
7008        while gi < gpr2 {
7009            let off0 = row_off + gi * TILE;
7010            let off1 = off0 + TILE;
7011            let s0 = f16_to_f32(u16::from_le_bytes([
7012                *bytes_ptr.add(off0),
7013                *bytes_ptr.add(off0 + 1),
7014            ]));
7015            let s1 = f16_to_f32(u16::from_le_bytes([
7016                *bytes_ptr.add(off1),
7017                *bytes_ptr.add(off1 + 1),
7018            ]));
7019            let (u0_0, u1_0, u2_0, u3_0) = q1t_unpack_reg_u64s(bytes_ptr.add(off0 + 2));
7020            let (u0_1, u1_1, u2_1, u3_1) = q1t_unpack_reg_u64s(bytes_ptr.add(off1 + 2));
7021            let w0_0 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u0_0), vcreate_u64(u1_0)));
7022            let w1_0 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u2_0), vcreate_u64(u3_0)));
7023            let w0_1 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u0_1), vcreate_u64(u1_1)));
7024            let w1_1 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u2_1), vcreate_u64(u3_1)));
7025            for k in 0..2 {
7026                let d0 = sdot2!(w0_0, w1_0, xp[k].add(gi * GROUP_SIZE));
7027                let d1 = sdot2!(w0_1, w1_1, xp[k].add((gi + 1) * GROUP_SIZE));
7028                acc[k] += d0 as f32 * s0 + d1 as f32 * s1;
7029            }
7030            gi += 2;
7031        }
7032        if gi < gpr {
7033            let off = row_off + gi * TILE;
7034            let s = f16_to_f32(u16::from_le_bytes([
7035                *bytes_ptr.add(off),
7036                *bytes_ptr.add(off + 1),
7037            ]));
7038            let (u0, u1, u2, u3) = q1t_unpack_reg_u64s(bytes_ptr.add(off + 2));
7039            let w0 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u0), vcreate_u64(u1)));
7040            let w1 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u2), vcreate_u64(u3)));
7041            for k in 0..2 {
7042                let d = sdot2!(w0, w1, xp[k].add(gi * GROUP_SIZE));
7043                acc[k] += d as f32 * s;
7044            }
7045        }
7046        acc
7047    }
7048}
7049
7050/// Fused Q1T pair matvec: ONE pass over the rows serves both
7051/// activation streams — on ARM the ternary register unpack happens
7052/// once per tile pair (`q1t_dot_row_sdot2`); elsewhere the second dot
7053/// rides the row's L1-warm tile bytes. Per stream the math matches
7054/// `q1t_matvec` exactly.
7055fn q1t_matvec2(
7056    bytes: &[u8],
7057    x1: &[f32],
7058    x2: &[f32],
7059    rows: usize,
7060    cols: usize,
7061    o1: &mut [f32],
7062    o2: &mut [f32],
7063    pool: Option<&Pool>,
7064) {
7065    debug_assert_eq!(o1.len(), rows);
7066    debug_assert_eq!(o2.len(), rows);
7067    const TILE: usize = cortiq_core::quant::Q1T_TILE;
7068    let gpr = cols / GROUP_SIZE;
7069    let (rp_off, ent_off, has_ov) = q1t_overlay(bytes, rows * gpr * TILE, rows);
7070    let out1 = SendMut(o1.as_mut_ptr());
7071    let out2 = SendMut(o2.as_mut_ptr());
7072    if a8w8_enabled() {
7073        let a1 = split_act(x1);
7074        let a2 = split_act(x2);
7075        let (a1, a2) = (&a1, &a2);
7076        let run = move |start: usize, end: usize| {
7077            for r in start..end {
7078                #[cfg(target_arch = "aarch64")]
7079                // a8w8 on aarch64 ⇔ sdot_enabled(), so the kernel's
7080                // target features are present.
7081                let ds = unsafe { q1t_dot_row_sdot2(bytes, r, gpr, &a1.xq, &a2.xq) };
7082                #[cfg(not(target_arch = "aarch64"))]
7083                let ds = [
7084                    q1t_dot_row_i8(bytes, r, gpr, &a1.xq),
7085                    q1t_dot_row_i8(bytes, r, gpr, &a2.xq),
7086                ];
7087                let mut acc1 = ds[0] * a1.sx;
7088                for &(j, xv) in &a1.outliers {
7089                    acc1 += q1t_base_weight(bytes, r, gpr, j) * xv;
7090                }
7091                acc1 += q1t_row_outlier_correction(bytes, r, rp_off, ent_off, has_ov, x1);
7092                let mut acc2 = ds[1] * a2.sx;
7093                for &(j, xv) in &a2.outliers {
7094                    acc2 += q1t_base_weight(bytes, r, gpr, j) * xv;
7095                }
7096                acc2 += q1t_row_outlier_correction(bytes, r, rp_off, ent_off, has_ov, x2);
7097                // SAFETY: disjoint row ranges per worker.
7098                unsafe {
7099                    *out1.at(r) = acc1;
7100                    *out2.at(r) = acc2;
7101                }
7102            }
7103        };
7104        dispatch_rows(pool, rows, &run);
7105        return;
7106    }
7107    let run = move |start: usize, end: usize| {
7108        // Exact path (CMF_SDOT=0): unpack the sign LUT once per group,
7109        // dot both streams — same op order per stream as `q1t_matvec`.
7110        let mut sg = [0f32; GROUP_SIZE];
7111        for r in start..end {
7112            let mut acc1 = 0f32;
7113            let mut acc2 = 0f32;
7114            for g in 0..gpr {
7115                let off = (r * gpr + g) * TILE;
7116                let s = f16_to_f32(u16::from_le_bytes([bytes[off], bytes[off + 1]]));
7117                let codes = &bytes[off + 2..off + TILE];
7118                for bi in 0..6 {
7119                    sg[bi * 5..bi * 5 + 5].copy_from_slice(&SIGN5[codes[bi] as usize]);
7120                }
7121                let lut = &SIGN5[codes[6] as usize];
7122                sg[30] = lut[0];
7123                sg[31] = lut[1];
7124                let xg1 = &x1[g * GROUP_SIZE..g * GROUP_SIZE + GROUP_SIZE];
7125                let xg2 = &x2[g * GROUP_SIZE..g * GROUP_SIZE + GROUP_SIZE];
7126                let mut gsum1 = 0f32;
7127                for k in 0..GROUP_SIZE {
7128                    gsum1 += sg[k] * xg1[k];
7129                }
7130                acc1 += s * gsum1;
7131                let mut gsum2 = 0f32;
7132                for k in 0..GROUP_SIZE {
7133                    gsum2 += sg[k] * xg2[k];
7134                }
7135                acc2 += s * gsum2;
7136            }
7137            acc1 += q1t_row_outlier_correction(bytes, r, rp_off, ent_off, has_ov, x1);
7138            acc2 += q1t_row_outlier_correction(bytes, r, rp_off, ent_off, has_ov, x2);
7139            // SAFETY: disjoint row ranges per worker.
7140            unsafe {
7141                *out1.at(r) = acc1;
7142                *out2.at(r) = acc2;
7143            }
7144        }
7145    };
7146    dispatch_rows(pool, rows, &run);
7147}
7148
7149/// Ternary (q1t) matmat (prefill) — dequant each row once, dot the whole
7150/// batch against it (amortizes the per-row decode).
7151fn q1t_matmat(
7152    bytes: &[u8],
7153    xs: &[f32],
7154    b: usize,
7155    rows: usize,
7156    cols: usize,
7157    out: &mut [f32],
7158    pool: Option<&Pool>,
7159) {
7160    debug_assert_eq!(out.len(), b * rows);
7161    const TILE: usize = cortiq_core::quant::Q1T_TILE;
7162    let gpr = cols / GROUP_SIZE;
7163    let (rp_off, ent_off, has_ov) = q1t_overlay(bytes, rows * gpr * TILE, rows);
7164    let out_addr = SendMut(out.as_mut_ptr());
7165    // int8 prefill (ARM SDOT / x86 AVX2): quantize the B inputs once, unpack
7166    // each weight row's signs to i8 ONCE, then int8-dot against every input —
7167    // the row sign-decode amortizes over the whole batch. CMF_SDOT=0 → f32.
7168    if a8w8_enabled() {
7169        let acts: Vec<SplitAct> = (0..b)
7170            .map(|bi| split_act(&xs[bi * cols..(bi + 1) * cols]))
7171            .collect();
7172        let acts = &acts;
7173        let run = move |start: usize, end: usize| {
7174            let mut sg = vec![0i8; cols + 8]; // row signs, i8 (+8 unpack slack)
7175            let mut sc = vec![0f32; gpr]; // per-group scales
7176            let mut accs = vec![0f32; b]; // per-batch accumulators, reused per row
7177            for r in start..end {
7178                for g in 0..gpr {
7179                    let off = (r * gpr + g) * TILE;
7180                    sc[g] = f16_to_f32(u16::from_le_bytes([bytes[off], bytes[off + 1]]));
7181                    q1t_unpack_group_i8(
7182                        bytes.as_ptr().wrapping_add(off + 2),
7183                        &mut sg[g * GROUP_SIZE..],
7184                    );
7185                }
7186                for bi in 0..b {
7187                    let act = &acts[bi];
7188                    let mut isum = 0f32;
7189                    for g in 0..gpr {
7190                        let d = q1t_i8dot32(
7191                            sg.as_ptr().wrapping_add(g * GROUP_SIZE),
7192                            act.xq.as_ptr().wrapping_add(g * GROUP_SIZE),
7193                        );
7194                        isum += d as f32 * sc[g];
7195                    }
7196                    let mut acc = isum * act.sx;
7197                    for &(j, xv) in &act.outliers {
7198                        acc += q1t_base_weight(bytes, r, gpr, j) * xv;
7199                    }
7200                    accs[bi] = acc;
7201                }
7202                // Overlay ONCE per row for the whole batch: read each (col, val)
7203                // from mmap a single time (was b× — the re-read dominated prefill)
7204                // and fan it out over the batch via the cached inputs.
7205                if has_ov {
7206                    let (c0, c1) = (
7207                        q1t_rowptr(bytes, rp_off, r),
7208                        q1t_rowptr(bytes, rp_off, r + 1),
7209                    );
7210                    for p in c0..c1 {
7211                        let e = ent_off + p * 4;
7212                        let col = u16::from_le_bytes([bytes[e], bytes[e + 1]]) as usize;
7213                        let val = f16_to_f32(u16::from_le_bytes([bytes[e + 2], bytes[e + 3]]));
7214                        for bi in 0..b {
7215                            accs[bi] += val * xs[bi * cols + col];
7216                        }
7217                    }
7218                }
7219                for bi in 0..b {
7220                    unsafe { *out_addr.at(bi * rows + r) = accs[bi] };
7221                }
7222            }
7223        };
7224        dispatch_rows(pool, rows, &run);
7225        return;
7226    }
7227    let run = move |start: usize, end: usize| {
7228        let mut buf = vec![0f32; cols];
7229        for r in start..end {
7230            q1t_dequant_row(bytes, r, gpr, rp_off, ent_off, has_ov, &mut buf);
7231            for bi in 0..b {
7232                let xr = &xs[bi * cols..(bi + 1) * cols];
7233                let mut acc = 0f32;
7234                for j in 0..cols {
7235                    acc += buf[j] * xr[j];
7236                }
7237                unsafe { *out_addr.at(bi * rows + r) = acc };
7238            }
7239        }
7240    };
7241    dispatch_rows(pool, rows, &run);
7242}
7243
7244fn q1_matvec(
7245    bytes: &[u8],
7246    x: &[f32],
7247    rows: usize,
7248    cols: usize,
7249    out: &mut [f32],
7250    pool: Option<&Pool>,
7251) {
7252    debug_assert_eq!(out.len(), rows);
7253    let gpr = cols / GROUP_SIZE;
7254    let out_addr = SendMut(out.as_mut_ptr());
7255    if a8w8_enabled() {
7256        let act = split_act(x);
7257        let gsum = q1_group_sums(&act.xq, gpr);
7258        let (act, gsum) = (&act, &gsum);
7259        let run = move |start: usize, end: usize| {
7260            q1_range_a8w8(bytes, gpr, act, gsum, out_addr, start, end)
7261        };
7262        dispatch_rows(pool, rows, &run);
7263        return;
7264    }
7265    let run = move |start: usize, end: usize| q1_range_f32(bytes, gpr, x, out_addr, start, end);
7266    dispatch_rows(pool, rows, &run);
7267}
7268
7269/// Fused two-input q1 matvec (weights read once per pair).
7270#[allow(clippy::too_many_arguments)]
7271fn q1_matvec2(
7272    bytes: &[u8],
7273    x1: &[f32],
7274    x2: &[f32],
7275    rows: usize,
7276    cols: usize,
7277    o1: &mut [f32],
7278    o2: &mut [f32],
7279    pool: Option<&Pool>,
7280) {
7281    let gpr = cols / GROUP_SIZE;
7282    let p1 = SendMut(o1.as_mut_ptr());
7283    let p2 = SendMut(o2.as_mut_ptr());
7284    if a8w8_enabled() {
7285        let a1 = split_act(x1);
7286        let a2 = split_act(x2);
7287        let g1 = q1_group_sums(&a1.xq, gpr);
7288        let g2 = q1_group_sums(&a2.xq, gpr);
7289        let (a1, a2, g1, g2) = (&a1, &a2, &g1, &g2);
7290        let run = move |start: usize, end: usize| {
7291            for r in start..end {
7292                let mut v1 = dot_q1_row_i8(bytes, r, gpr, &a1.xq, g1) * a1.sx;
7293                let mut v2 = dot_q1_row_i8(bytes, r, gpr, &a2.xq, g2) * a2.sx;
7294                for &(j, xv) in &a1.outliers {
7295                    let (w, s) = q1_outlier(bytes, r, gpr, j);
7296                    v1 += w * s * xv;
7297                }
7298                for &(j, xv) in &a2.outliers {
7299                    let (w, s) = q1_outlier(bytes, r, gpr, j);
7300                    v2 += w * s * xv;
7301                }
7302                // SAFETY: disjoint row ranges per worker.
7303                unsafe {
7304                    *p1.at(r) = v1;
7305                    *p2.at(r) = v2;
7306                }
7307            }
7308        };
7309        dispatch_rows(pool, rows, &run);
7310        return;
7311    }
7312    let run = move |start: usize, end: usize| {
7313        for r in start..end {
7314            // SAFETY: disjoint row ranges per worker.
7315            unsafe {
7316                *p1.at(r) = q1_row_exact(bytes, r, gpr, x1);
7317                *p2.at(r) = q1_row_exact(bytes, r, gpr, x2);
7318            }
7319        }
7320    };
7321    dispatch_rows(pool, rows, &run);
7322}
7323
7324/// Batched q1 matmat: each row's tiles stream once per microbatch.
7325#[allow(clippy::too_many_arguments)]
7326fn q1_matmat(
7327    bytes: &[u8],
7328    xs_all: &[f32],
7329    b: usize,
7330    rows: usize,
7331    cols: usize,
7332    out: &mut [f32],
7333    pool: Option<&Pool>,
7334) {
7335    debug_assert_eq!(out.len(), b * rows);
7336    let gpr = cols / GROUP_SIZE;
7337    let out_addr = SendMut(out.as_mut_ptr());
7338    if a8w8_enabled() {
7339        let acts: Vec<(SplitAct, Vec<i32>)> = (0..b)
7340            .map(|bi| {
7341                let act = split_act(&xs_all[bi * cols..(bi + 1) * cols]);
7342                let gsum = q1_group_sums(&act.xq, gpr);
7343                (act, gsum)
7344            })
7345            .collect();
7346        let acts = &acts;
7347        #[cfg(target_arch = "x86_64")]
7348        let blocked_ok = avx2_enabled() && blocked_enabled();
7349        #[cfg(target_arch = "aarch64")]
7350        let blocked_ok = sdot_enabled() && blocked_enabled();
7351        let run = move |start: usize, end: usize| {
7352            for r in start..end {
7353                let mut bi = 0usize;
7354                // Blocked 1×4: the unpacked bit mask serves four
7355                // activation streams per group.
7356                #[cfg(target_arch = "aarch64")]
7357                if blocked_ok {
7358                    while bi + 4 <= acts.len() {
7359                        let xs = [
7360                            acts[bi].0.xq.as_slice(),
7361                            acts[bi + 1].0.xq.as_slice(),
7362                            acts[bi + 2].0.xq.as_slice(),
7363                            acts[bi + 3].0.xq.as_slice(),
7364                        ];
7365                        let gs = [
7366                            acts[bi].1.as_slice(),
7367                            acts[bi + 1].1.as_slice(),
7368                            acts[bi + 2].1.as_slice(),
7369                            acts[bi + 3].1.as_slice(),
7370                        ];
7371                        let d = unsafe { dot_q1_row_1x4_sdot(bytes, r, gpr, xs, gs) };
7372                        for k in 0..4 {
7373                            let (act, _) = &acts[bi + k];
7374                            let mut acc = d[k] * act.sx;
7375                            for &(j, xv) in &act.outliers {
7376                                let (w, sc) = q1_outlier(bytes, r, gpr, j);
7377                                acc += w * sc * xv;
7378                            }
7379                            // SAFETY: disjoint (bi, r) cells per worker.
7380                            unsafe { *out_addr.at((bi + k) * rows + r) = acc };
7381                        }
7382                        bi += 4;
7383                    }
7384                }
7385                #[cfg(target_arch = "x86_64")]
7386                if blocked_ok {
7387                    while bi + 4 <= acts.len() {
7388                        let xs = [
7389                            acts[bi].0.xq.as_slice(),
7390                            acts[bi + 1].0.xq.as_slice(),
7391                            acts[bi + 2].0.xq.as_slice(),
7392                            acts[bi + 3].0.xq.as_slice(),
7393                        ];
7394                        let gs = [
7395                            acts[bi].1.as_slice(),
7396                            acts[bi + 1].1.as_slice(),
7397                            acts[bi + 2].1.as_slice(),
7398                            acts[bi + 3].1.as_slice(),
7399                        ];
7400                        let d = unsafe {
7401                            if vnni_tiles_enabled() {
7402                                dot_q1_row_1x4_vnni(bytes, r, gpr, xs, gs)
7403                            } else {
7404                                dot_q1_row_1x4_avx2(bytes, r, gpr, xs, gs)
7405                            }
7406                        };
7407                        for k in 0..4 {
7408                            let (act, _) = &acts[bi + k];
7409                            let mut acc = d[k] * act.sx;
7410                            for &(j, xv) in &act.outliers {
7411                                let (w, sc) = q1_outlier(bytes, r, gpr, j);
7412                                acc += w * sc * xv;
7413                            }
7414                            // SAFETY: disjoint (bi, r) cells per worker.
7415                            unsafe { *out_addr.at((bi + k) * rows + r) = acc };
7416                        }
7417                        bi += 4;
7418                    }
7419                }
7420                while bi < acts.len() {
7421                    let (act, gsum) = &acts[bi];
7422                    let mut acc = dot_q1_row_i8(bytes, r, gpr, &act.xq, gsum) * act.sx;
7423                    for &(j, xv) in &act.outliers {
7424                        let (w, s) = q1_outlier(bytes, r, gpr, j);
7425                        acc += w * s * xv;
7426                    }
7427                    // SAFETY: disjoint (bi, r) cells per worker range.
7428                    unsafe { *out_addr.at(bi * rows + r) = acc };
7429                    bi += 1;
7430                }
7431            }
7432        };
7433        dispatch_rows(pool, rows, &run);
7434        return;
7435    }
7436    let run = move |start: usize, end: usize| {
7437        for r in start..end {
7438            for bi in 0..b {
7439                let x = &xs_all[bi * cols..(bi + 1) * cols];
7440                // SAFETY: disjoint (bi, r) cells per worker range.
7441                unsafe { *out_addr.at(bi * rows + r) = q1_row_exact(bytes, r, gpr, x) };
7442            }
7443        }
7444    };
7445    dispatch_rows(pool, rows, &run);
7446}
7447
7448/// Fused q4_block matvec straight from the mapped bytes. SDOT path when
7449/// dotprod is available (port of vmfcore `dot_q4_block_sdot`, measured
7450/// +23% on q4 decode): nibbles → centered i8, int8×int8 `sdot` per
7451/// 32-group, exact outlier correction — the same A8W8 contract as q8.
7452/// `CMF_SDOT=0` keeps the exact scalar path.
7453fn q4matvec(
7454    bytes: &[u8],
7455    x: &[f32],
7456    rows: usize,
7457    cols: usize,
7458    out: &mut [f32],
7459    pool: Option<&Pool>,
7460) {
7461    debug_assert_eq!(out.len(), rows);
7462    let (packed, scales) = q4_split(bytes, rows, cols);
7463    let gpr = cols / GROUP_SIZE;
7464    let out_addr = SendMut(out.as_mut_ptr());
7465
7466    if a8w8_enabled() {
7467        let act = split_act(x);
7468        let run = move |start: usize, end: usize| {
7469            q4_range_a8w8(packed, scales, gpr, cols, &act, out_addr, start, end)
7470        };
7471        dispatch_rows(pool, rows, &run);
7472        return;
7473    }
7474
7475    let run =
7476        move |start: usize, end: usize| q4_range_f32(packed, scales, gpr, x, out_addr, start, end);
7477    dispatch_rows(pool, rows, &run);
7478}
7479
7480/// One q4 row via the A8W8 int8 path — SDOT on ARM, AVX2 maddubs on
7481/// x86 (scalar fallback is unreachable: callers gate on a8w8_enabled).
7482#[inline]
7483#[allow(unreachable_code)]
7484/// One UNPACKED q4 row (centered i8 in `buf`) against four activation
7485/// streams: the 32-byte weight chunk and its abs() load once per group,
7486/// the per-group f16 scale decodes once — four maddubs+reduce chains
7487/// instead of four full (load, abs, dot) rounds.
7488#[cfg(target_arch = "x86_64")]
7489#[target_feature(enable = "avx2")]
7490unsafe fn dot_q4b_row_1x4_avx2(
7491    buf: &[u8],
7492    scales: &[u8],
7493    g0: usize,
7494    gpr: usize,
7495    xs: [&[i8]; 4],
7496) -> [f32; 4] {
7497    // SAFETY: callers uphold buffer contracts (buf.len() == gpr·32).
7498    unsafe {
7499        use core::arch::x86_64::*;
7500        let ones = _mm256_set1_epi16(1);
7501        let mut acc = [0f32; 4];
7502        for gi in 0..gpr {
7503            let s = f16_to_f32(u16::from_le_bytes([
7504                scales[(g0 + gi) * 2],
7505                scales[(g0 + gi) * 2 + 1],
7506            ]));
7507            let w = _mm256_loadu_si256(buf.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
7508            let aw = _mm256_abs_epi8(w);
7509            for (k, xq) in xs.iter().enumerate() {
7510                let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
7511                let p16 = _mm256_maddubs_epi16(aw, _mm256_sign_epi8(x, w));
7512                let d = _mm256_madd_epi16(p16, ones);
7513                let hi128 = _mm256_extracti128_si256::<1>(d);
7514                let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
7515                let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
7516                let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
7517                acc[k] += _mm_cvtsi128_si32(s32) as f32 * s;
7518            }
7519        }
7520        acc
7521    }
7522}
7523
7524/// VNNI twin of `dot_q4b_row_1x4_avx2` (see `dpbusd_hsum`).
7525#[cfg(target_arch = "x86_64")]
7526#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
7527unsafe fn dot_q4b_row_1x4_vnni(
7528    buf: &[u8],
7529    scales: &[u8],
7530    g0: usize,
7531    gpr: usize,
7532    xs: [&[i8]; 4],
7533) -> [f32; 4] {
7534    // SAFETY: callers uphold buffer contracts (buf.len() == gpr·32).
7535    unsafe {
7536        use core::arch::x86_64::*;
7537        let mut acc = [0f32; 4];
7538        for gi in 0..gpr {
7539            let s = f16_to_f32(u16::from_le_bytes([
7540                scales[(g0 + gi) * 2],
7541                scales[(g0 + gi) * 2 + 1],
7542            ]));
7543            let w = _mm256_loadu_si256(buf.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
7544            let aw = _mm256_abs_epi8(w);
7545            for (k, xq) in xs.iter().enumerate() {
7546                let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
7547                let d = dpbusd_hsum(aw, _mm256_sign_epi8(x, w));
7548                acc[k] += d as f32 * s;
7549            }
7550        }
7551        acc
7552    }
7553}
7554
7555/// The vbit flavor of the blocked 1×4: the per-activation A8W8 scale
7556/// folds in PER GROUP as `(d·sx)·s` — bit-matching the single-matvec
7557/// accumulation order (the q4_block flavor applies sx once at the end,
7558/// matching ITS single path; the two conventions are historical and
7559/// each blocked leg must mirror its own).
7560#[cfg(target_arch = "x86_64")]
7561#[target_feature(enable = "avx2")]
7562unsafe fn dot_q4b_row_1x4_sx_avx2(
7563    buf: &[u8],
7564    scales: &[u8],
7565    g0: usize,
7566    gpr: usize,
7567    xs: [&[i8]; 4],
7568    sxs: [f32; 4],
7569) -> [f32; 4] {
7570    // SAFETY: callers uphold buffer contracts (buf.len() == gpr·32).
7571    unsafe {
7572        use core::arch::x86_64::*;
7573        let ones = _mm256_set1_epi16(1);
7574        let mut acc = [0f32; 4];
7575        for gi in 0..gpr {
7576            let s = f16_to_f32(u16::from_le_bytes([
7577                scales[(g0 + gi) * 2],
7578                scales[(g0 + gi) * 2 + 1],
7579            ]));
7580            let w = _mm256_loadu_si256(buf.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
7581            let aw = _mm256_abs_epi8(w);
7582            for (k, xq) in xs.iter().enumerate() {
7583                let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
7584                let p16 = _mm256_maddubs_epi16(aw, _mm256_sign_epi8(x, w));
7585                let d = _mm256_madd_epi16(p16, ones);
7586                let hi128 = _mm256_extracti128_si256::<1>(d);
7587                let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
7588                let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
7589                let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
7590                acc[k] += (_mm_cvtsi128_si32(s32) as f32 * sxs[k]) * s;
7591            }
7592        }
7593        acc
7594    }
7595}
7596
7597/// VNNI twin of `dot_q4b_row_1x4_sx_avx2` (see `dpbusd_hsum`; the
7598/// per-group `(d·sx)·s` fold mirrors the vbit single path).
7599#[cfg(target_arch = "x86_64")]
7600#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
7601unsafe fn dot_q4b_row_1x4_sx_vnni(
7602    buf: &[u8],
7603    scales: &[u8],
7604    g0: usize,
7605    gpr: usize,
7606    xs: [&[i8]; 4],
7607    sxs: [f32; 4],
7608) -> [f32; 4] {
7609    // SAFETY: callers uphold buffer contracts (buf.len() == gpr·32).
7610    unsafe {
7611        use core::arch::x86_64::*;
7612        let mut acc = [0f32; 4];
7613        for gi in 0..gpr {
7614            let s = f16_to_f32(u16::from_le_bytes([
7615                scales[(g0 + gi) * 2],
7616                scales[(g0 + gi) * 2 + 1],
7617            ]));
7618            let w = _mm256_loadu_si256(buf.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
7619            let aw = _mm256_abs_epi8(w);
7620            for (k, xq) in xs.iter().enumerate() {
7621                let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
7622                let d = dpbusd_hsum(aw, _mm256_sign_epi8(x, w));
7623                acc[k] += (d as f32 * sxs[k]) * s;
7624            }
7625        }
7626        acc
7627    }
7628}
7629
7630#[allow(unreachable_code)]
7631fn dot_q4_row_i8(packed: &[u8], scales: &[u8], g0: usize, gpr: usize, xq: &[i8]) -> f32 {
7632    #[cfg(target_arch = "aarch64")]
7633    unsafe {
7634        return dot_q4_row_sdot(packed, scales, g0, gpr, xq);
7635    }
7636    #[cfg(target_arch = "x86_64")]
7637    unsafe {
7638        return dot_q4_row_avx2(packed, scales, g0, gpr, xq);
7639    }
7640    let mut acc = 0f32;
7641    for gi in 0..gpr {
7642        let g = g0 + gi;
7643        let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
7644        let mut d = 0i32;
7645        for (k, &b) in packed[g * 16..(g + 1) * 16].iter().enumerate() {
7646            d += ((b & 0x0F) as i32 - 8) * xq[gi * GROUP_SIZE + k * 2] as i32
7647                + (((b >> 4) & 0x0F) as i32 - 8) * xq[gi * GROUP_SIZE + k * 2 + 1] as i32;
7648        }
7649        acc += d as f32 * s;
7650    }
7651    acc
7652}
7653
7654/// Two-activation q4 row via the A8W8 int8 path (see `dot_q4_row_i8`).
7655#[inline]
7656#[allow(unreachable_code)]
7657fn dot_q4_row_i8_2(
7658    packed: &[u8],
7659    scales: &[u8],
7660    g0: usize,
7661    gpr: usize,
7662    xq1: &[i8],
7663    xq2: &[i8],
7664) -> (f32, f32) {
7665    #[cfg(target_arch = "aarch64")]
7666    unsafe {
7667        return dot_q4_row_sdot2(packed, scales, g0, gpr, xq1, xq2);
7668    }
7669    #[cfg(target_arch = "x86_64")]
7670    unsafe {
7671        return dot_q4_row_avx2_2(packed, scales, g0, gpr, xq1, xq2);
7672    }
7673    (
7674        dot_q4_row_i8(packed, scales, g0, gpr, xq1),
7675        dot_q4_row_i8(packed, scales, g0, gpr, xq2),
7676    )
7677}
7678
7679/// One q4 row range via SDOT (kernel body of `q4matvec`, extracted so
7680/// multi-matrix jobs can drive it for several tensors in one dispatch).
7681#[allow(clippy::too_many_arguments)]
7682fn q4_range_a8w8(
7683    packed: &[u8],
7684    scales: &[u8],
7685    gpr: usize,
7686    cols: usize,
7687    act: &SplitAct,
7688    out: SendMut,
7689    start: usize,
7690    end: usize,
7691) {
7692    for r in start..end {
7693        let mut acc = dot_q4_row_i8(packed, scales, r * gpr, gpr, &act.xq) * act.sx;
7694        // xq is zeroed at outlier slots — add the exact terms.
7695        for &(j, xv) in &act.outliers {
7696            let flat = r * cols + j;
7697            let byte = packed[flat / 2];
7698            let nib = if flat & 1 == 0 {
7699                byte & 0x0F
7700            } else {
7701                byte >> 4
7702            };
7703            let s = f16_to_f32(u16::from_le_bytes([
7704                scales[(flat / GROUP_SIZE) * 2],
7705                scales[(flat / GROUP_SIZE) * 2 + 1],
7706            ]));
7707            acc += ((nib as i32 - 8) as f32) * s * xv;
7708        }
7709        // SAFETY: disjoint row ranges per worker.
7710        unsafe { *out.at(r) = acc };
7711    }
7712}
7713
7714/// Two-input q4 row range via the A8W8 int8 path — kernel body of
7715/// `q4matvec2`, extracted for pair multi-matrix jobs.
7716#[allow(clippy::too_many_arguments)]
7717fn q4_range2_a8w8(
7718    packed: &[u8],
7719    scales: &[u8],
7720    gpr: usize,
7721    cols: usize,
7722    a1: &SplitAct,
7723    a2: &SplitAct,
7724    p1: SendMut,
7725    p2: SendMut,
7726    start: usize,
7727    end: usize,
7728) {
7729    for r in start..end {
7730        let (s1, s2) = dot_q4_row_i8_2(packed, scales, r * gpr, gpr, &a1.xq, &a2.xq);
7731        let mut acc1 = s1 * a1.sx;
7732        let mut acc2 = s2 * a2.sx;
7733        // xq is zeroed at outlier slots — add the exact terms.
7734        let fix = |outliers: &[(usize, f32)], acc: &mut f32| {
7735            for &(j, xv) in outliers {
7736                let flat = r * cols + j;
7737                let byte = packed[flat / 2];
7738                let nib = if flat & 1 == 0 {
7739                    byte & 0x0F
7740                } else {
7741                    byte >> 4
7742                };
7743                let s = f16_to_f32(u16::from_le_bytes([
7744                    scales[(flat / GROUP_SIZE) * 2],
7745                    scales[(flat / GROUP_SIZE) * 2 + 1],
7746                ]));
7747                *acc += ((nib as i32 - 8) as f32) * s * xv;
7748            }
7749        };
7750        fix(&a1.outliers, &mut acc1);
7751        fix(&a2.outliers, &mut acc2);
7752        // SAFETY: disjoint row ranges per worker.
7753        unsafe {
7754            *p1.at(r) = acc1;
7755            *p2.at(r) = acc2;
7756        }
7757    }
7758}
7759
7760/// Exact scalar q4 row range (same extraction, non-SDOT path).
7761fn q4_range_f32(
7762    packed: &[u8],
7763    scales: &[u8],
7764    gpr: usize,
7765    x: &[f32],
7766    out: SendMut,
7767    start: usize,
7768    end: usize,
7769) {
7770    for r in start..end {
7771        let mut acc = 0f32;
7772        for gi in 0..gpr {
7773            let g = r * gpr + gi;
7774            let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
7775            let pk = &packed[g * 16..(g + 1) * 16];
7776            let xg = &x[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
7777            let mut ga = 0f32;
7778            for (k, &b) in pk.iter().enumerate() {
7779                ga += ((b & 0x0F) as f32 - 8.0) * xg[k * 2]
7780                    + (((b >> 4) & 0x0F) as f32 - 8.0) * xg[k * 2 + 1];
7781            }
7782            acc += ga * s;
7783        }
7784        // SAFETY: disjoint row ranges per worker.
7785        unsafe { *out.at(r) = acc };
7786    }
7787}
7788
7789/// Fused two-input q4 matvec: nibbles are unpacked ONCE per group and
7790/// dotted against both activations (was: two full matvecs — double
7791/// weight traffic). Per-lane math matches `q4matvec` exactly.
7792#[allow(clippy::too_many_arguments)]
7793fn q4matvec2(
7794    bytes: &[u8],
7795    x1: &[f32],
7796    x2: &[f32],
7797    rows: usize,
7798    cols: usize,
7799    o1: &mut [f32],
7800    o2: &mut [f32],
7801    pool: Option<&Pool>,
7802) {
7803    debug_assert_eq!(o1.len(), rows);
7804    debug_assert_eq!(o2.len(), rows);
7805    let (packed, scales) = q4_split(bytes, rows, cols);
7806    let gpr = cols / GROUP_SIZE;
7807
7808    if a8w8_enabled() {
7809        let a1 = split_act(x1);
7810        let a2 = split_act(x2);
7811        let p1 = SendMut(o1.as_mut_ptr());
7812        let p2 = SendMut(o2.as_mut_ptr());
7813        let run = move |start: usize, end: usize| {
7814            q4_range2_a8w8(packed, scales, gpr, cols, &a1, &a2, p1, p2, start, end)
7815        };
7816        dispatch_rows(pool, rows, &run);
7817        return;
7818    }
7819
7820    let p1 = SendMut(o1.as_mut_ptr());
7821    let p2 = SendMut(o2.as_mut_ptr());
7822    let run = move |start: usize, end: usize| {
7823        q4_range2_f32(packed, scales, gpr, x1, x2, p1, p2, start, end)
7824    };
7825    dispatch_rows(pool, rows, &run);
7826}
7827
7828/// Two-input exact scalar q4 row range (same extraction).
7829#[allow(clippy::too_many_arguments)]
7830fn q4_range2_f32(
7831    packed: &[u8],
7832    scales: &[u8],
7833    gpr: usize,
7834    x1: &[f32],
7835    x2: &[f32],
7836    p1: SendMut,
7837    p2: SendMut,
7838    start: usize,
7839    end: usize,
7840) {
7841    for r in start..end {
7842        let (mut acc1, mut acc2) = (0f32, 0f32);
7843        for gi in 0..gpr {
7844            let g = r * gpr + gi;
7845            let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
7846            let pk = &packed[g * 16..(g + 1) * 16];
7847            let x1g = &x1[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
7848            let x2g = &x2[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
7849            let (mut g1, mut g2) = (0f32, 0f32);
7850            for (k, &b) in pk.iter().enumerate() {
7851                let wl = (b & 0x0F) as f32 - 8.0;
7852                let wh = ((b >> 4) & 0x0F) as f32 - 8.0;
7853                g1 += wl * x1g[k * 2] + wh * x1g[k * 2 + 1];
7854                g2 += wl * x2g[k * 2] + wh * x2g[k * 2 + 1];
7855            }
7856            acc1 += g1 * s;
7857            acc2 += g2 * s;
7858        }
7859        // SAFETY: disjoint row ranges per worker.
7860        unsafe {
7861            *p1.at(r) = acc1;
7862            *p2.at(r) = acc2;
7863        }
7864    }
7865}
7866
7867thread_local! {
7868    /// Per-worker decoded-row scratch for the batched q4/vbit kernels
7869    /// (centered i8 for SDOT, f32 for the exact/scalar paths).
7870    static ROW_I8: std::cell::RefCell<Vec<u8>> = const { std::cell::RefCell::new(Vec::new()) };
7871    static ROW_F32: std::cell::RefCell<Vec<f32>> = const { std::cell::RefCell::new(Vec::new()) };
7872}
7873
7874/// Batched q4 matmat: each weight row is unpacked from the mmap ONCE
7875/// and dotted against ALL b activations (prefill used to fall back to b
7876/// full matvecs — b× weight traffic and b× nibble decode). Per-position
7877/// math matches `q4matvec` exactly: same group order, same accumulation.
7878/// `out` is row-major [b, rows] like `qmatmat`.
7879#[allow(clippy::too_many_arguments)]
7880fn q4matmat(
7881    bytes: &[u8],
7882    xs_all: &[f32],
7883    b: usize,
7884    rows: usize,
7885    cols: usize,
7886    out: &mut [f32],
7887    pool: Option<&Pool>,
7888) {
7889    debug_assert_eq!(xs_all.len(), b * cols);
7890    debug_assert_eq!(out.len(), b * rows);
7891    let (packed, scales) = q4_split(bytes, rows, cols);
7892    let gpr = cols / GROUP_SIZE;
7893    let gscale = |g: usize| f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
7894
7895    if a8w8_enabled() {
7896        let acts: Vec<SplitAct> = (0..b)
7897            .map(|bi| split_act(&xs_all[bi * cols..(bi + 1) * cols]))
7898            .collect();
7899        let acts = &acts;
7900        let out_addr = SendMut(out.as_mut_ptr());
7901        let run = move |start: usize, end: usize| {
7902            ROW_I8.with(|rb| {
7903                let mut buf = rb.borrow_mut();
7904                buf.resize(cols, 0);
7905                for r in start..end {
7906                    // Unpack the row's nibbles to centered i8 once
7907                    // (element 2k = low nibble, 2k+1 = high — flat order,
7908                    // same as dot_q4_row_sdot's zip).
7909                    for gi in 0..gpr {
7910                        let g = r * gpr + gi;
7911                        for (k, &bt) in packed[g * 16..(g + 1) * 16].iter().enumerate() {
7912                            buf[gi * GROUP_SIZE + k * 2] = ((bt & 0x0F) as i32 - 8) as i8 as u8;
7913                            buf[gi * GROUP_SIZE + k * 2 + 1] =
7914                                (((bt >> 4) & 0x0F) as i32 - 8) as i8 as u8;
7915                        }
7916                    }
7917                    let mut bi = 0usize;
7918                    #[cfg(target_arch = "x86_64")]
7919                    if avx2_enabled() && blocked_enabled() {
7920                        while bi + 4 <= acts.len() {
7921                            let xs = [
7922                                acts[bi].xq.as_slice(),
7923                                acts[bi + 1].xq.as_slice(),
7924                                acts[bi + 2].xq.as_slice(),
7925                                acts[bi + 3].xq.as_slice(),
7926                            ];
7927                            let d = unsafe {
7928                                if vnni_tiles_enabled() {
7929                                    dot_q4b_row_1x4_vnni(&buf, scales, r * gpr, gpr, xs)
7930                                } else {
7931                                    dot_q4b_row_1x4_avx2(&buf, scales, r * gpr, gpr, xs)
7932                                }
7933                            };
7934                            for k in 0..4 {
7935                                let act = &acts[bi + k];
7936                                let mut acc = d[k] * act.sx;
7937                                for &(j, xv) in &act.outliers {
7938                                    acc += (buf[j] as i8) as f32
7939                                        * gscale((r * cols + j) / GROUP_SIZE)
7940                                        * xv;
7941                                }
7942                                // SAFETY: disjoint (bi, r) cells per worker.
7943                                unsafe { *out_addr.at((bi + k) * rows + r) = acc };
7944                            }
7945                            bi += 4;
7946                        }
7947                    }
7948                    while bi < acts.len() {
7949                        let act = &acts[bi];
7950                        let mut acc = 0f32;
7951                        for gi in 0..gpr {
7952                            let d = dot_i8_i8(
7953                                &buf[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE],
7954                                &act.xq[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE],
7955                            );
7956                            acc += d as f32 * gscale(r * gpr + gi);
7957                        }
7958                        acc *= act.sx;
7959                        // xq is zeroed at outlier slots — exact terms.
7960                        for &(j, xv) in &act.outliers {
7961                            acc += (buf[j] as i8) as f32 * gscale((r * cols + j) / GROUP_SIZE) * xv;
7962                        }
7963                        // SAFETY: disjoint (bi, r) cells per worker row range.
7964                        unsafe { *out_addr.at(bi * rows + r) = acc };
7965                        bi += 1;
7966                    }
7967                }
7968            })
7969        };
7970        dispatch_rows(pool, rows, &run);
7971        return;
7972    }
7973
7974    let out_addr = SendMut(out.as_mut_ptr());
7975    let run = move |start: usize, end: usize| {
7976        ROW_F32.with(|rb| {
7977            let mut buf = rb.borrow_mut();
7978            buf.resize(cols, 0.0);
7979            for r in start..end {
7980                // Decode raw (nib − 8) values once; scales stay per-group
7981                // so the accumulation order matches q4matvec bit-for-bit.
7982                for gi in 0..gpr {
7983                    let g = r * gpr + gi;
7984                    for (k, &bt) in packed[g * 16..(g + 1) * 16].iter().enumerate() {
7985                        buf[gi * GROUP_SIZE + k * 2] = (bt & 0x0F) as f32 - 8.0;
7986                        buf[gi * GROUP_SIZE + k * 2 + 1] = ((bt >> 4) & 0x0F) as f32 - 8.0;
7987                    }
7988                }
7989                for bi in 0..b {
7990                    let x = &xs_all[bi * cols..(bi + 1) * cols];
7991                    let mut acc = 0f32;
7992                    for gi in 0..gpr {
7993                        let mut ga = 0f32;
7994                        // Pairwise (lo + hi) addition, matching
7995                        // q4matvec's `ga += lo·x + hi·x` shape exactly —
7996                        // a flat one-per-element loop rounds differently
7997                        // and broke bit-parity on the scalar (x86) path.
7998                        for k in 0..GROUP_SIZE / 2 {
7999                            let e = gi * GROUP_SIZE + k * 2;
8000                            ga += buf[e] * x[e] + buf[e + 1] * x[e + 1];
8001                        }
8002                        acc += ga * gscale(r * gpr + gi);
8003                    }
8004                    // SAFETY: disjoint (bi, r) cells per worker row range.
8005                    unsafe { *out_addr.at(bi * rows + r) = acc };
8006                }
8007            }
8008        })
8009    };
8010    dispatch_rows(pool, rows, &run);
8011}
8012
8013/// Batched vbit matmat: each variable-bit row is decoded from the mmap
8014/// ONCE for the whole microbatch. Same per-position math as
8015/// `vbitmatvec` (SDOT A8W8 with exact outliers / exact f32 for b=8 rows
8016/// and the scalar path).
8017#[allow(clippy::too_many_arguments)]
8018fn vbitmatmat(
8019    bytes: &[u8],
8020    offsets: &[usize],
8021    xs_all: &[f32],
8022    b: usize,
8023    rows: usize,
8024    cols: usize,
8025    out: &mut [f32],
8026    pool: Option<&Pool>,
8027) {
8028    debug_assert_eq!(xs_all.len(), b * cols);
8029    debug_assert_eq!(out.len(), b * rows);
8030    debug_assert_eq!(offsets.len(), rows + 1);
8031    let ng = cols / GROUP_SIZE;
8032    let bits = &bytes[..rows];
8033    let sc_off = rows;
8034    let gscale = |r: usize, g: usize| {
8035        let so = (r * ng + g) * 2;
8036        f16_to_f32(u16::from_le_bytes([
8037            bytes[sc_off + so],
8038            bytes[sc_off + so + 1],
8039        ]))
8040    };
8041
8042    // Decode row r's raw (u − L) values into `dst` (f32, unscaled).
8043    let decode_f32 = |r: usize, dst: &mut [f32]| {
8044        let bw = bits[r] as usize;
8045        let l = ((1i32 << (bw - 1)) - 1) as f32;
8046        let data = &bytes[offsets[r]..offsets[r + 1]];
8047        let (mut acc, mut nbits, mut idx) = (0u64, 0usize, 0usize);
8048        for d in dst.iter_mut() {
8049            while nbits < bw {
8050                acc = (acc << 8) | data[idx] as u64;
8051                idx += 1;
8052                nbits += 8;
8053            }
8054            let u = ((acc >> (nbits - bw)) & ((1u64 << bw) - 1)) as f32;
8055            nbits -= bw;
8056            *d = u - l;
8057        }
8058    };
8059
8060    if a8w8_enabled() {
8061        let acts: Vec<SplitAct> = (0..b)
8062            .map(|bi| split_act(&xs_all[bi * cols..(bi + 1) * cols]))
8063            .collect();
8064        let acts = &acts;
8065        let out_addr = SendMut(out.as_mut_ptr());
8066        let run = move |start: usize, end: usize| {
8067            for r in start..end {
8068                let bw = bits[r] as usize;
8069                if bw == 8 {
8070                    // u−L reaches 128 → no i8 path; decode once, exact
8071                    // f32 dots for every position (same as vbitmatvec).
8072                    ROW_F32.with(|rb| {
8073                        let mut buf = rb.borrow_mut();
8074                        buf.resize(cols, 0.0);
8075                        decode_f32(r, &mut buf);
8076                        for bi in 0..b {
8077                            let x = &xs_all[bi * cols..(bi + 1) * cols];
8078                            let mut dot = 0f32;
8079                            for g in 0..ng {
8080                                let mut gd = 0f32;
8081                                for k in 0..GROUP_SIZE {
8082                                    gd += buf[g * GROUP_SIZE + k] * x[g * GROUP_SIZE + k];
8083                                }
8084                                dot += gd * gscale(r, g);
8085                            }
8086                            // SAFETY: disjoint (bi, r) cells per worker range.
8087                            unsafe { *out_addr.at(bi * rows + r) = dot };
8088                        }
8089                    });
8090                    continue;
8091                }
8092                let l = (1i32 << (bw - 1)) - 1;
8093                let data = &bytes[offsets[r]..offsets[r + 1]];
8094                ROW_I8.with(|rb| {
8095                    let mut buf = rb.borrow_mut();
8096                    buf.resize(cols, 0);
8097                    #[inline(always)]
8098                    fn fill<const B: usize>(data: &[u8], l: i32, buf: &mut [u8]) {
8099                        for (blk, chunk) in buf.chunks_exact_mut(8).enumerate() {
8100                            let u = unpack8::<B>(&data[blk * B..]);
8101                            for k in 0..8 {
8102                                chunk[k] = (u[k] - l) as i8 as u8;
8103                            }
8104                        }
8105                    }
8106                    match bw {
8107                        3 => fill::<3>(data, l, &mut buf),
8108                        4 => vbit_fill4(data, &mut buf),
8109                        5 => fill::<5>(data, l, &mut buf),
8110                        6 => fill::<6>(data, l, &mut buf),
8111                        _ => unreachable!("vbit bit-width {bw} (validated at load)"),
8112                    }
8113                    let mut bi = 0usize;
8114                    // The vbit scale table shares q4_block's layout
8115                    // (contiguous f16 per (row·ng + g)), so the same
8116                    // blocked 1×4 kernel serves the decoded row.
8117                    #[cfg(target_arch = "x86_64")]
8118                    if avx2_enabled() && blocked_enabled() {
8119                        while bi + 4 <= acts.len() {
8120                            let xs = [
8121                                acts[bi].xq.as_slice(),
8122                                acts[bi + 1].xq.as_slice(),
8123                                acts[bi + 2].xq.as_slice(),
8124                                acts[bi + 3].xq.as_slice(),
8125                            ];
8126                            let sxs = [
8127                                acts[bi].sx,
8128                                acts[bi + 1].sx,
8129                                acts[bi + 2].sx,
8130                                acts[bi + 3].sx,
8131                            ];
8132                            let d = unsafe {
8133                                if vnni_tiles_enabled() {
8134                                    dot_q4b_row_1x4_sx_vnni(
8135                                        &buf,
8136                                        &bytes[sc_off..],
8137                                        r * ng,
8138                                        ng,
8139                                        xs,
8140                                        sxs,
8141                                    )
8142                                } else {
8143                                    dot_q4b_row_1x4_sx_avx2(
8144                                        &buf,
8145                                        &bytes[sc_off..],
8146                                        r * ng,
8147                                        ng,
8148                                        xs,
8149                                        sxs,
8150                                    )
8151                                }
8152                            };
8153                            for k in 0..4 {
8154                                let act = &acts[bi + k];
8155                                let mut dot = d[k];
8156                                for &(j, xv) in &act.outliers {
8157                                    dot += (buf[j] as i8) as f32 * gscale(r, j / GROUP_SIZE) * xv;
8158                                }
8159                                // SAFETY: disjoint (bi, r) cells per worker.
8160                                unsafe { *out_addr.at((bi + k) * rows + r) = dot };
8161                            }
8162                            bi += 4;
8163                        }
8164                    }
8165                    while bi < acts.len() {
8166                        let act = &acts[bi];
8167                        let mut dot = 0f32;
8168                        for g in 0..ng {
8169                            let d = dot_i8_i8(
8170                                &buf[g * GROUP_SIZE..(g + 1) * GROUP_SIZE],
8171                                &act.xq[g * GROUP_SIZE..(g + 1) * GROUP_SIZE],
8172                            ) as f32
8173                                * act.sx;
8174                            dot += d * gscale(r, g);
8175                        }
8176                        for &(j, xv) in &act.outliers {
8177                            dot += (buf[j] as i8) as f32 * gscale(r, j / GROUP_SIZE) * xv;
8178                        }
8179                        // SAFETY: disjoint (bi, r) cells per worker range.
8180                        unsafe { *out_addr.at(bi * rows + r) = dot };
8181                        bi += 1;
8182                    }
8183                });
8184            }
8185        };
8186        dispatch_rows(pool, rows, &run);
8187        return;
8188    }
8189
8190    let out_addr = SendMut(out.as_mut_ptr());
8191    let run = move |start: usize, end: usize| {
8192        ROW_F32.with(|rb| {
8193            let mut buf = rb.borrow_mut();
8194            buf.resize(cols, 0.0);
8195            for r in start..end {
8196                decode_f32(r, &mut buf);
8197                for bi in 0..b {
8198                    let x = &xs_all[bi * cols..(bi + 1) * cols];
8199                    let mut dot = 0f32;
8200                    for g in 0..ng {
8201                        let mut gd = 0f32;
8202                        for k in 0..GROUP_SIZE {
8203                            gd += buf[g * GROUP_SIZE + k] * x[g * GROUP_SIZE + k];
8204                        }
8205                        dot += gd * gscale(r, g);
8206                    }
8207                    // SAFETY: disjoint (bi, r) cells per worker range.
8208                    unsafe { *out_addr.at(bi * rows + r) = dot };
8209                }
8210            }
8211        })
8212    };
8213    dispatch_rows(pool, rows, &run);
8214}
8215
8216/// Build a GPU batch job for a q8-family mapped tensor (primary
8217/// shard): prescaled input + directory coordinates. None → not
8218/// GPU-eligible, caller stays on the CPU.
8219pub(crate) fn gpu_batch_job<'a>(
8220    t: &'a QTensor,
8221    x: &[f32],
8222) -> Option<(std::sync::Arc<CmfModel>, crate::gpu::BatchJob<'a>)> {
8223    match t {
8224        QTensor::Mapped {
8225            model,
8226            idx,
8227            dtype: dt @ (TensorDtype::Q8Row | TensorDtype::Q8_2f),
8228            rows,
8229            cols,
8230            row_scale,
8231            col_field,
8232            ..
8233        } => Some((
8234            model.clone(),
8235            crate::gpu::BatchJob {
8236                idx: *idx,
8237                rows: *rows,
8238                cols: *cols,
8239                row_scale,
8240                xs: prescale(x, col_field, *dt).into_owned(),
8241                layout: crate::gpu::BatchLayout::Q8,
8242            },
8243        )),
8244        // q1: raw f32 activations, tile-embedded scales.
8245        QTensor::Mapped {
8246            model,
8247            idx,
8248            dtype: TensorDtype::Q1,
8249            rows,
8250            cols,
8251            ..
8252        } => Some((
8253            model.clone(),
8254            crate::gpu::BatchJob {
8255                idx: *idx,
8256                rows: *rows,
8257                cols: *cols,
8258                row_scale: &[],
8259                xs: x.to_vec(),
8260                layout: crate::gpu::BatchLayout::Q1,
8261            },
8262        )),
8263        // q4_tiled / q4tp: raw f32 activations; the scales live in the
8264        // payload (inline tiles / row ladder), so row_scale stays empty.
8265        // The GDN projection batch already runs these layouts on Metal —
8266        // this arm lets the attention QKV batch reach the same kernels.
8267        QTensor::Mapped {
8268            model,
8269            idx,
8270            dtype: dt @ (TensorDtype::Q4Tiled | TensorDtype::Q4TiledP),
8271            rows,
8272            cols,
8273            ..
8274        } => Some((
8275            model.clone(),
8276            crate::gpu::BatchJob {
8277                idx: *idx,
8278                rows: *rows,
8279                cols: *cols,
8280                row_scale: &[],
8281                xs: x.to_vec(),
8282                layout: if *dt == TensorDtype::Q4Tiled {
8283                    crate::gpu::BatchLayout::Q4t
8284                } else {
8285                    crate::gpu::BatchLayout::Q4tp
8286                },
8287            },
8288        )),
8289        _ => None,
8290    }
8291}
8292
8293thread_local! {
8294    static PRESCALE_BUF1: std::cell::RefCell<Vec<f32>> = const { std::cell::RefCell::new(Vec::new()) };
8295    static PRESCALE_BUF2: std::cell::RefCell<Vec<f32>> = const { std::cell::RefCell::new(Vec::new()) };
8296}
8297
8298pub(crate) fn prescale<'a>(
8299    x: &'a [f32],
8300    col_field: &[f32],
8301    dtype: TensorDtype,
8302) -> std::borrow::Cow<'a, [f32]> {
8303    if dtype == TensorDtype::Q8_2f {
8304        x.iter().zip(col_field).map(|(a, c)| a * c).collect()
8305    } else {
8306        std::borrow::Cow::Borrowed(x)
8307    }
8308}
8309
8310/// θ col-field fold for q8_2f activations. Borrowed pass-through for
8311/// every other dtype, using thread-local buffers to eliminate per-matvec allocations.
8312pub(crate) fn prescale_with<R, F: FnOnce(&[f32]) -> R>(
8313    x: &[f32],
8314    col_field: &[f32],
8315    dtype: TensorDtype,
8316    buf_id: u8,
8317    f: F,
8318) -> R {
8319    if dtype == TensorDtype::Q8_2f {
8320        if buf_id == 1 {
8321            PRESCALE_BUF1.with(|b| {
8322                let mut buf = b.borrow_mut();
8323                buf.clear();
8324                buf.extend(x.iter().zip(col_field).map(|(a, c)| a * c));
8325                f(&buf)
8326            })
8327        } else {
8328            PRESCALE_BUF2.with(|b| {
8329                let mut buf = b.borrow_mut();
8330                buf.clear();
8331                buf.extend(x.iter().zip(col_field).map(|(a, c)| a * c));
8332                f(&buf)
8333            })
8334        }
8335    } else {
8336        f(x)
8337    }
8338}
8339
8340// ───────────────────── x86-64 AVX2 kernels (roadmap этап 2) ─────────────────────
8341
8342/// AVX2+FMA available? Default ON when the CPU supports both;
8343/// `CMF_AVX2=0` disables (falls back to the autovectorized loops).
8344#[cfg(target_arch = "x86_64")]
8345pub(crate) fn avx2_enabled() -> bool {
8346    use std::sync::OnceLock;
8347    static ON: OnceLock<bool> = OnceLock::new();
8348    *ON.get_or_init(|| {
8349        std::env::var("CMF_AVX2").map(|v| v != "0").unwrap_or(true)
8350            && std::arch::is_x86_feature_detected!("avx2")
8351            && std::arch::is_x86_feature_detected!("fma")
8352    })
8353}
8354
8355/// AVX2 A8W8 allowed? The quantized-activation contract is switched by
8356/// the SAME env as the ARM SDOT path: `CMF_SDOT=0` keeps exact kernels
8357/// (the golden-parity exact gate relies on it) — AVX2 f32 kernels stay
8358/// active either way, they are exact (regrouped sums only).
8359#[cfg(target_arch = "x86_64")]
8360fn avx2_a8w8_enabled() -> bool {
8361    use std::sync::OnceLock;
8362    static ON: OnceLock<bool> = OnceLock::new();
8363    *ON.get_or_init(|| {
8364        avx2_enabled() && std::env::var("CMF_SDOT").map(|v| v != "0").unwrap_or(true)
8365    })
8366}
8367
8368/// A8W8 quantized-activation path available on THIS machine? One
8369/// switch across architectures: ARM dotprod (CMF_SDOT) or x86 AVX2
8370/// (CMF_AVX2 + the same CMF_SDOT exact-contract override).
8371#[inline]
8372pub(crate) fn a8w8_enabled() -> bool {
8373    #[cfg(target_arch = "aarch64")]
8374    {
8375        sdot_enabled()
8376    }
8377    #[cfg(target_arch = "x86_64")]
8378    {
8379        avx2_a8w8_enabled()
8380    }
8381    #[cfg(not(any(target_arch = "aarch64", target_arch = "x86_64")))]
8382    {
8383        false
8384    }
8385}
8386
8387/// int8·int8 dot dispatch: SDOT on ARM; AVX-512 VNNI (vpdpbusd) or AVX2
8388/// maddubs on x86. Callers are gated by `a8w8_enabled()`.
8389#[inline]
8390#[allow(unreachable_code)]
8391fn dot_i8_i8(w: &[u8], xq: &[i8]) -> i32 {
8392    #[cfg(target_arch = "aarch64")]
8393    unsafe {
8394        return dot_i8_sdot(w, xq);
8395    }
8396    #[cfg(target_arch = "x86_64")]
8397    unsafe {
8398        if avx512vnni_enabled() {
8399            return dot_i8_i8_vnni(w, xq);
8400        }
8401        return dot_i8_i8_avx2(w, xq);
8402    }
8403    w.iter()
8404        .zip(xq)
8405        .map(|(&a, &b)| (a as i8) as i32 * b as i32)
8406        .sum()
8407}
8408
8409/// AVX-512 VNNI available? (F+BW+VL+VNNI; `CMF_AVX512=0` falls back to
8410/// AVX2.) VL matters: short 32-byte groups (q4/vbit) ride the 256-bit
8411/// `vpdpbusd` encoding.
8412#[cfg(target_arch = "x86_64")]
8413fn avx512vnni_enabled() -> bool {
8414    use std::sync::OnceLock;
8415    static ON: OnceLock<bool> = OnceLock::new();
8416    *ON.get_or_init(|| {
8417        std::env::var("CMF_AVX512")
8418            .map(|v| v != "0")
8419            .unwrap_or(true)
8420            && std::arch::is_x86_feature_detected!("avx512f")
8421            && std::arch::is_x86_feature_detected!("avx512bw")
8422            && std::arch::is_x86_feature_detected!("avx512vl")
8423            && std::arch::is_x86_feature_detected!("avx512vnni")
8424    })
8425}
8426
8427/// Grouped-codec VNNI arms (the q4t/q4b/q1/q1t tile kernels): default
8428/// ON where AVX-512 VNNI exists (`CMF_VNNI_TILES=0` opt-out). Measured
8429/// on Ryzen 7950X (Zen4, 3 alternating process pairs, blocked GEMM
8430/// 4864×896 b=256): q4t 63→68 GF/s (+8%), q1 53→56 (+6%), q4b 72→75
8431/// (+4%) — consistent, no leg regressed. The tile kernels keep a
8432/// horizontal reduce per 32-weight group, so the `vpdpbusd` saving is
8433/// smaller than the long-dot q8 win (+13%), but it is real and free.
8434#[cfg(target_arch = "x86_64")]
8435fn vnni_tiles_enabled() -> bool {
8436    use std::sync::OnceLock;
8437    static ON: OnceLock<bool> = OnceLock::new();
8438    *ON.get_or_init(|| {
8439        std::env::var("CMF_VNNI_TILES")
8440            .map(|v| v != "0")
8441            .unwrap_or(true)
8442            && avx512vnni_enabled()
8443    })
8444}
8445
8446/// One 256-bit u8×i8 dot → i32 via `vpdpbusd` into a fresh accumulator
8447/// plus the same horizontal reduce the AVX2 kernels use. Products are
8448/// bounded (|w| ≤ 8 or ≤ 1), so maddubs never saturated — the i32 sum
8449/// is bit-identical to the maddubs+madd pair it replaces.
8450#[cfg(target_arch = "x86_64")]
8451#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
8452#[inline]
8453unsafe fn dpbusd_hsum(aw: core::arch::x86_64::__m256i, xs: core::arch::x86_64::__m256i) -> i32 {
8454    // SAFETY: pure register math.
8455    unsafe {
8456        use core::arch::x86_64::*;
8457        let d = _mm256_dpbusd_epi32(_mm256_setzero_si256(), aw, xs);
8458        let hi128 = _mm256_extracti128_si256::<1>(d);
8459        let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
8460        let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
8461        let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
8462        _mm_cvtsi128_si32(s32)
8463    }
8464}
8465
8466/// int8·int8 via AVX-512 VNNI: `vpdpbusd` fuses the maddubs+madd+add
8467/// triple into one u8×i8 dot-accumulate. AVX-512 has no vpsignb, so the
8468/// |w|·sign(x,w) trick becomes |w| × (x negated where w<0) via a mask
8469/// subtract — w==0 lanes contribute 0 through |w|=0 either way.
8470#[cfg(target_arch = "x86_64")]
8471#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
8472unsafe fn dot_i8_i8_vnni(w: &[u8], xq: &[i8]) -> i32 {
8473    // SAFETY: callers uphold slice-length contracts (see call sites).
8474    unsafe {
8475        use core::arch::x86_64::*;
8476        let n = w.len();
8477        let mut j = 0usize;
8478        let mut total: i32;
8479        // 4 independent accumulators: vpdpbusd is its own loop-carried
8480        // dependency (~5-cycle latency) — a single-acc loop runs
8481        // latency-bound and LOSES to the AVX2 maddubs kernel, measured
8482        // on Granite Rapids.
8483        {
8484            #[inline(always)]
8485            unsafe fn step(
8486                w: *const u8,
8487                x: *const i8,
8488                acc: core::arch::x86_64::__m512i,
8489            ) -> core::arch::x86_64::__m512i {
8490                unsafe {
8491                    use core::arch::x86_64::*;
8492                    let wv = _mm512_loadu_si512(w as *const _);
8493                    let xv = _mm512_loadu_si512(x as *const _);
8494                    let aw = _mm512_abs_epi8(wv);
8495                    let neg = _mm512_movepi8_mask(wv);
8496                    let sx = _mm512_mask_sub_epi8(xv, neg, _mm512_setzero_si512(), xv);
8497                    _mm512_dpbusd_epi32(acc, aw, sx)
8498                }
8499            }
8500            let (mut a0, mut a1, mut a2, mut a3) = (
8501                _mm512_setzero_si512(),
8502                _mm512_setzero_si512(),
8503                _mm512_setzero_si512(),
8504                _mm512_setzero_si512(),
8505            );
8506            while j + 256 <= n {
8507                a0 = step(w.as_ptr().add(j), xq.as_ptr().add(j), a0);
8508                a1 = step(w.as_ptr().add(j + 64), xq.as_ptr().add(j + 64), a1);
8509                a2 = step(w.as_ptr().add(j + 128), xq.as_ptr().add(j + 128), a2);
8510                a3 = step(w.as_ptr().add(j + 192), xq.as_ptr().add(j + 192), a3);
8511                j += 256;
8512            }
8513            while j + 64 <= n {
8514                a0 = step(w.as_ptr().add(j), xq.as_ptr().add(j), a0);
8515                j += 64;
8516            }
8517            let s01 = _mm512_add_epi32(a0, a1);
8518            let s23 = _mm512_add_epi32(a2, a3);
8519            total = _mm512_reduce_add_epi32(_mm512_add_epi32(s01, s23));
8520        }
8521        // 32-wide (q4/vbit groups are exactly 32 bytes).
8522        if j + 32 <= n {
8523            let wv = _mm256_loadu_si256(w.as_ptr().add(j) as *const __m256i);
8524            let xv = _mm256_loadu_si256(xq.as_ptr().add(j) as *const __m256i);
8525            let d = _mm256_dpbusd_epi32(
8526                _mm256_setzero_si256(),
8527                _mm256_abs_epi8(wv),
8528                _mm256_sign_epi8(xv, wv),
8529            );
8530            let hi128 = _mm256_extracti128_si256::<1>(d);
8531            let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
8532            let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
8533            let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
8534            total += _mm_cvtsi128_si32(s32);
8535            j += 32;
8536        }
8537        while j < n {
8538            total += (w[j] as i8) as i32 * xq[j] as i32;
8539            j += 1;
8540        }
8541        total
8542    }
8543}
8544
8545/// i8 row · f32 x via AVX2/FMA (x86 mirror of `dot_i8_f32_neon`).
8546#[cfg(target_arch = "x86_64")]
8547#[target_feature(enable = "avx2,fma")]
8548unsafe fn dot_i8_f32_avx2(w: &[u8], x: &[f32]) -> f32 {
8549    // SAFETY: callers uphold slice-length contracts (see call sites).
8550    unsafe {
8551        use core::arch::x86_64::*;
8552        let n = x.len();
8553        let wp = w.as_ptr();
8554        let xp = x.as_ptr();
8555        let (mut a0, mut a1) = (_mm256_setzero_ps(), _mm256_setzero_ps());
8556        let mut j = 0usize;
8557        while j + 16 <= n {
8558            let wb = _mm_loadu_si128(wp.add(j) as *const __m128i);
8559            let lo = _mm256_cvtepi8_epi32(wb);
8560            let hi = _mm256_cvtepi8_epi32(_mm_srli_si128::<8>(wb));
8561            a0 = _mm256_fmadd_ps(_mm256_cvtepi32_ps(lo), _mm256_loadu_ps(xp.add(j)), a0);
8562            a1 = _mm256_fmadd_ps(_mm256_cvtepi32_ps(hi), _mm256_loadu_ps(xp.add(j + 8)), a1);
8563            j += 16;
8564        }
8565        let acc = _mm256_add_ps(a0, a1);
8566        let hi128 = _mm256_extractf128_ps::<1>(acc);
8567        let s128 = _mm_add_ps(_mm256_castps256_ps128(acc), hi128);
8568        let s64 = _mm_add_ps(s128, _mm_movehl_ps(s128, s128));
8569        let s32 = _mm_add_ss(s64, _mm_shuffle_ps::<1>(s64, s64));
8570        let mut sum = _mm_cvtss_f32(s32);
8571        while j < n {
8572            sum += (*wp.add(j) as i8) as f32 * *xp.add(j);
8573            j += 1;
8574        }
8575        sum
8576    }
8577}
8578
8579/// int8(weight)·int8(activation) → i32 via AVX2 maddubs — the x86
8580/// analogue of the SDOT A8W8 path. `maddubs` takes u8×i8, so the
8581/// standard sign trick applies: |w| × sign(x, w) ≡ w × x per lane.
8582/// Pair saturation is safe: |w|≤128, |x|≤127 → 2·128·127 < 32767.
8583#[cfg(target_arch = "x86_64")]
8584#[target_feature(enable = "avx2")]
8585unsafe fn dot_i8_i8_avx2(w: &[u8], xq: &[i8]) -> i32 {
8586    // SAFETY: callers uphold slice-length contracts (see call sites).
8587    unsafe {
8588        use core::arch::x86_64::*;
8589        let n = w.len();
8590        let ones = _mm256_set1_epi16(1);
8591        let mut acc = _mm256_setzero_si256();
8592        let mut j = 0usize;
8593        while j + 32 <= n {
8594            let wv = _mm256_loadu_si256(w.as_ptr().add(j) as *const __m256i);
8595            let xv = _mm256_loadu_si256(xq.as_ptr().add(j) as *const __m256i);
8596            let p16 = _mm256_maddubs_epi16(_mm256_abs_epi8(wv), _mm256_sign_epi8(xv, wv));
8597            acc = _mm256_add_epi32(acc, _mm256_madd_epi16(p16, ones));
8598            j += 32;
8599        }
8600        let hi128 = _mm256_extracti128_si256::<1>(acc);
8601        let s128 = _mm_add_epi32(_mm256_castsi256_si128(acc), hi128);
8602        let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
8603        let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
8604        let mut s = _mm_cvtsi128_si32(s32);
8605        while j < n {
8606            s += (w[j] as i8) as i32 * xq[j] as i32;
8607            j += 1;
8608        }
8609        s
8610    }
8611}
8612
8613/// smmla 2×4: one instruction covers a 2-row × 2-activation × 8-deep
8614/// tile (32 MACs vs sdot's 16) — the weight pair loads once per 8-k
8615/// slice as a combined 2×8 register and meets two activation pairs.
8616#[cfg(target_arch = "aarch64")]
8617#[target_feature(enable = "neon,i8mm")]
8618unsafe fn dot_i8_smmla_2x4(w0: &[u8], w1: &[u8], xs: [&[i8]; 4]) -> [[i32; 4]; 2] {
8619    // SAFETY: callers uphold slice-length contracts.
8620    unsafe {
8621        use core::arch::aarch64::*;
8622        use core::arch::asm;
8623        let n = w0.len();
8624        let w0p = w0.as_ptr() as *const i8;
8625        let w1p = w1.as_ptr() as *const i8;
8626        // acc01 holds [c(r0,x0) c(r0,x1) c(r1,x0) c(r1,x1)]; acc23 the
8627        // same for x2/x3.
8628        let mut acc01 = vdupq_n_s32(0);
8629        let mut acc23 = vdupq_n_s32(0);
8630        let mut i = 0usize;
8631        while i + 8 <= n {
8632            let wa = vcombine_s8(vld1_s8(w0p.add(i)), vld1_s8(w1p.add(i)));
8633            let xb01 = vcombine_s8(
8634                vld1_s8(xs[0].as_ptr().add(i)),
8635                vld1_s8(xs[1].as_ptr().add(i)),
8636            );
8637            let xb23 = vcombine_s8(
8638                vld1_s8(xs[2].as_ptr().add(i)),
8639                vld1_s8(xs[3].as_ptr().add(i)),
8640            );
8641            asm!(
8642                "smmla {a01:v}.4s, {w:v}.16b, {x01:v}.16b",
8643                "smmla {a23:v}.4s, {w:v}.16b, {x23:v}.16b",
8644                a01 = inout(vreg) acc01, a23 = inout(vreg) acc23,
8645                w = in(vreg) wa, x01 = in(vreg) xb01, x23 = in(vreg) xb23,
8646                options(pure, nomem, nostack),
8647            );
8648            i += 8;
8649        }
8650        let mut out = [[0i32; 4]; 2];
8651        let a01: [i32; 4] = core::mem::transmute(acc01);
8652        let a23: [i32; 4] = core::mem::transmute(acc23);
8653        out[0][0] = a01[0];
8654        out[0][1] = a01[1];
8655        out[1][0] = a01[2];
8656        out[1][1] = a01[3];
8657        out[0][2] = a23[0];
8658        out[0][3] = a23[1];
8659        out[1][2] = a23[2];
8660        out[1][3] = a23[3];
8661        if i < n {
8662            for (k, x) in xs.iter().enumerate() {
8663                for j in i..n {
8664                    out[0][k] += (w0[j] as i8) as i32 * x[j] as i32;
8665                    out[1][k] += (w1[j] as i8) as i32 * x[j] as i32;
8666                }
8667            }
8668        }
8669        out
8670    }
8671}
8672
8673/// ARM twin of the x86 blocked prefill GEMM: two weight rows stay in
8674/// registers across four activation streams, eight sdot accumulators.
8675/// (The per-row form re-read each W row once per activation.)
8676#[cfg(target_arch = "aarch64")]
8677#[target_feature(enable = "neon,dotprod")]
8678unsafe fn dot_i8_sdot_2x4(w0: &[u8], w1: &[u8], xs: [&[i8]; 4]) -> [[i32; 4]; 2] {
8679    // SAFETY: callers uphold slice-length contracts.
8680    unsafe {
8681        use core::arch::aarch64::*;
8682        use core::arch::asm;
8683        let n = w0.len();
8684        let w0p = w0.as_ptr() as *const i8;
8685        let w1p = w1.as_ptr() as *const i8;
8686        let mut acc = [[vdupq_n_s32(0); 4]; 2];
8687        let mut i = 0usize;
8688        while i + 16 <= n {
8689            let wv0 = vld1q_s8(w0p.add(i));
8690            let wv1 = vld1q_s8(w1p.add(i));
8691            for (k, x) in xs.iter().enumerate() {
8692                let xv = vld1q_s8(x.as_ptr().add(i));
8693                let (mut a0, mut a1) = (acc[0][k], acc[1][k]);
8694                asm!(
8695                    "sdot {a0:v}.4s, {w0:v}.16b, {x:v}.16b",
8696                    "sdot {a1:v}.4s, {w1:v}.16b, {x:v}.16b",
8697                    a0 = inout(vreg) a0, a1 = inout(vreg) a1,
8698                    w0 = in(vreg) wv0, w1 = in(vreg) wv1, x = in(vreg) xv,
8699                    options(pure, nomem, nostack),
8700                );
8701                acc[0][k] = a0;
8702                acc[1][k] = a1;
8703            }
8704            i += 16;
8705        }
8706        let mut out = [[0i32; 4]; 2];
8707        for r in 0..2 {
8708            for k in 0..4 {
8709                out[r][k] = vaddvq_s32(acc[r][k]);
8710            }
8711        }
8712        if i < n {
8713            for (k, x) in xs.iter().enumerate() {
8714                for j in i..n {
8715                    out[0][k] += (w0[j] as i8) as i32 * x[j] as i32;
8716                    out[1][k] += (w1[j] as i8) as i32 * x[j] as i32;
8717                }
8718            }
8719        }
8720        out
8721    }
8722}
8723
8724/// Blocked 2 weight rows × 4 activations for the prefill GEMM
8725/// (roadmap P0: packed panels + multi-row accumulators). The two rows'
8726/// abs() live in registers across all four activation streams; the
8727/// sign-fixup is recomputed per pair (the price of the maddubs trick).
8728/// Returns raw i8·i8 dots; the caller applies scales and outliers.
8729#[cfg(target_arch = "x86_64")]
8730#[target_feature(enable = "avx2")]
8731unsafe fn dot_i8_i8_avx2_2x4(w0: &[u8], w1: &[u8], xs: [&[i8]; 4]) -> [[i32; 4]; 2] {
8732    // SAFETY: callers uphold slice-length contracts.
8733    unsafe {
8734        use core::arch::x86_64::*;
8735        let n = w0.len();
8736        let ones = _mm256_set1_epi16(1);
8737        let mut acc = [[_mm256_setzero_si256(); 4]; 2];
8738        let mut j = 0usize;
8739        while j + 32 <= n {
8740            let wv0 = _mm256_loadu_si256(w0.as_ptr().add(j) as *const __m256i);
8741            let wv1 = _mm256_loadu_si256(w1.as_ptr().add(j) as *const __m256i);
8742            let aw0 = _mm256_abs_epi8(wv0);
8743            let aw1 = _mm256_abs_epi8(wv1);
8744            for (k, x) in xs.iter().enumerate() {
8745                let xv = _mm256_loadu_si256(x.as_ptr().add(j) as *const __m256i);
8746                let p0 = _mm256_maddubs_epi16(aw0, _mm256_sign_epi8(xv, wv0));
8747                acc[0][k] = _mm256_add_epi32(acc[0][k], _mm256_madd_epi16(p0, ones));
8748                let p1 = _mm256_maddubs_epi16(aw1, _mm256_sign_epi8(xv, wv1));
8749                acc[1][k] = _mm256_add_epi32(acc[1][k], _mm256_madd_epi16(p1, ones));
8750            }
8751            j += 32;
8752        }
8753        let mut out = [[0i32; 4]; 2];
8754        for r in 0..2 {
8755            for k in 0..4 {
8756                let a = acc[r][k];
8757                let hi128 = _mm256_extracti128_si256::<1>(a);
8758                let s128 = _mm_add_epi32(_mm256_castsi256_si128(a), hi128);
8759                let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
8760                let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
8761                out[r][k] = _mm_cvtsi128_si32(s32);
8762            }
8763        }
8764        if j < n {
8765            for (k, x) in xs.iter().enumerate() {
8766                for i in j..n {
8767                    out[0][k] += (w0[i] as i8) as i32 * x[i] as i32;
8768                    out[1][k] += (w1[i] as i8) as i32 * x[i] as i32;
8769                }
8770            }
8771        }
8772        out
8773    }
8774}
8775
8776/// AVX2/VNNI q8 row dot with exact outlier correction (x86 mirror of
8777/// `row_dot_sdot` — same A8W8 contract). With AVX-512 VNNI the row goes
8778/// through the bias trick: Σ(w+128)·x via pure `vpdpbusd` (no per-lane
8779/// sign fixups), corrected by −128·Σx with Σx precomputed per split.
8780#[cfg(target_arch = "x86_64")]
8781#[inline]
8782fn row_dot_avx2(row: &[u8], act: &SplitAct) -> f32 {
8783    let dot = if avx512vnni_enabled() && row.len() >= 64 {
8784        (unsafe { dot_u8p128_i8_vnni(row, &act.xq) }) - 128 * act.xsum
8785    } else {
8786        unsafe { dot_i8_i8_avx2(row, &act.xq) }
8787    };
8788    let mut acc = dot as f32 * act.sx;
8789    for &(j, xv) in &act.outliers {
8790        acc += (row[j] as i8) as f32 * xv;
8791    }
8792    acc
8793}
8794
8795/// Σ (w[i]+128)·x[i] via pure `vpdpbusd` — the caller subtracts
8796/// 128·Σx. Four independent accumulators (dpbusd is ~5-cycle latency;
8797/// a single-acc loop runs latency-bound, measured on Granite Rapids).
8798#[cfg(target_arch = "x86_64")]
8799#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
8800unsafe fn dot_u8p128_i8_vnni(w: &[u8], xq: &[i8]) -> i32 {
8801    // SAFETY: callers uphold slice-length contracts (see call sites).
8802    unsafe {
8803        use core::arch::x86_64::*;
8804        let n = w.len();
8805        let flip = _mm512_set1_epi8(-128); // XOR 0x80: i8 w → u8 (w+128)
8806        #[inline(always)]
8807        unsafe fn step(
8808            w: *const u8,
8809            x: *const i8,
8810            flip: core::arch::x86_64::__m512i,
8811            acc: core::arch::x86_64::__m512i,
8812        ) -> core::arch::x86_64::__m512i {
8813            unsafe {
8814                use core::arch::x86_64::*;
8815                let wv = _mm512_xor_si512(_mm512_loadu_si512(w as *const _), flip);
8816                _mm512_dpbusd_epi32(acc, wv, _mm512_loadu_si512(x as *const _))
8817            }
8818        }
8819        let (mut a0, mut a1, mut a2, mut a3) = (
8820            _mm512_setzero_si512(),
8821            _mm512_setzero_si512(),
8822            _mm512_setzero_si512(),
8823            _mm512_setzero_si512(),
8824        );
8825        let mut j = 0usize;
8826        while j + 256 <= n {
8827            a0 = step(w.as_ptr().add(j), xq.as_ptr().add(j), flip, a0);
8828            a1 = step(w.as_ptr().add(j + 64), xq.as_ptr().add(j + 64), flip, a1);
8829            a2 = step(w.as_ptr().add(j + 128), xq.as_ptr().add(j + 128), flip, a2);
8830            a3 = step(w.as_ptr().add(j + 192), xq.as_ptr().add(j + 192), flip, a3);
8831            j += 256;
8832        }
8833        while j + 64 <= n {
8834            a0 = step(w.as_ptr().add(j), xq.as_ptr().add(j), flip, a0);
8835            j += 64;
8836        }
8837        let mut total = _mm512_reduce_add_epi32(_mm512_add_epi32(
8838            _mm512_add_epi32(a0, a1),
8839            _mm512_add_epi32(a2, a3),
8840        ));
8841        // Scalar tail: (w as i8) + 128 ≡ (w as u8) ^ 0x80.
8842        while j < n {
8843            total += ((w[j] ^ 0x80) as i32) * xq[j] as i32;
8844            j += 1;
8845        }
8846        total
8847    }
8848}
8849
8850/// One q4 row via AVX2: nibbles → centered i8 (unpacklo/hi restores the
8851/// writer's flat order, same as the NEON vzip pair), maddubs against
8852/// the pre-quantized activation group, × the group's f16 scale. Pair
8853/// saturation safe: |w|≤8, |x|≤127 → 2·8·127 ≪ 32767. Mirror of
8854/// `dot_q4_row_sdot`.
8855#[cfg(target_arch = "x86_64")]
8856#[target_feature(enable = "avx2")]
8857unsafe fn dot_q4_row_avx2(packed: &[u8], scales: &[u8], g0: usize, gpr: usize, xq: &[i8]) -> f32 {
8858    // SAFETY: callers uphold slice-length contracts (16 packed bytes and
8859    // 2 scale bytes per group; xq.len() == gpr·GROUP_SIZE).
8860    unsafe {
8861        use core::arch::x86_64::*;
8862        let lomask = _mm_set1_epi8(0x0F);
8863        let eight = _mm256_set1_epi8(8);
8864        let ones = _mm256_set1_epi16(1);
8865        let mut acc = 0f32;
8866        for gi in 0..gpr {
8867            let g = g0 + gi;
8868            let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
8869            let b = _mm_loadu_si128(packed.as_ptr().add(g * 16) as *const __m128i);
8870            let lo = _mm_and_si128(b, lomask);
8871            let hi = _mm_and_si128(_mm_srli_epi16::<4>(b), lomask);
8872            let w = _mm256_sub_epi8(
8873                _mm256_set_m128i(_mm_unpackhi_epi8(lo, hi), _mm_unpacklo_epi8(lo, hi)),
8874                eight,
8875            );
8876            let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
8877            let p16 = _mm256_maddubs_epi16(_mm256_abs_epi8(w), _mm256_sign_epi8(x, w));
8878            let d = _mm256_madd_epi16(p16, ones);
8879            let hi128 = _mm256_extracti128_si256::<1>(d);
8880            let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
8881            let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
8882            let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
8883            acc += _mm_cvtsi128_si32(s32) as f32 * s;
8884        }
8885        acc
8886    }
8887}
8888
8889/// Two-activation q4 row via AVX2: nibbles unpacked ONCE per group,
8890/// both activations dotted against the same centered i8 register.
8891#[cfg(target_arch = "x86_64")]
8892#[target_feature(enable = "avx2")]
8893unsafe fn dot_q4_row_avx2_2(
8894    packed: &[u8],
8895    scales: &[u8],
8896    g0: usize,
8897    gpr: usize,
8898    xq1: &[i8],
8899    xq2: &[i8],
8900) -> (f32, f32) {
8901    // SAFETY: callers uphold slice-length contracts (see dot_q4_row_avx2).
8902    unsafe {
8903        use core::arch::x86_64::*;
8904        let lomask = _mm_set1_epi8(0x0F);
8905        let eight = _mm256_set1_epi8(8);
8906        let ones = _mm256_set1_epi16(1);
8907        let (mut acc1, mut acc2) = (0f32, 0f32);
8908        #[inline(always)]
8909        unsafe fn hsum(d: core::arch::x86_64::__m256i) -> i32 {
8910            unsafe {
8911                use core::arch::x86_64::*;
8912                let hi128 = _mm256_extracti128_si256::<1>(d);
8913                let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
8914                let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
8915                let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
8916                _mm_cvtsi128_si32(s32)
8917            }
8918        }
8919        for gi in 0..gpr {
8920            let g = g0 + gi;
8921            let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
8922            let b = _mm_loadu_si128(packed.as_ptr().add(g * 16) as *const __m128i);
8923            let lo = _mm_and_si128(b, lomask);
8924            let hi = _mm_and_si128(_mm_srli_epi16::<4>(b), lomask);
8925            let w = _mm256_sub_epi8(
8926                _mm256_set_m128i(_mm_unpackhi_epi8(lo, hi), _mm_unpacklo_epi8(lo, hi)),
8927                eight,
8928            );
8929            let aw = _mm256_abs_epi8(w);
8930            let x1 = _mm256_loadu_si256(xq1.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
8931            let x2 = _mm256_loadu_si256(xq2.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
8932            let d1 = _mm256_madd_epi16(_mm256_maddubs_epi16(aw, _mm256_sign_epi8(x1, w)), ones);
8933            let d2 = _mm256_madd_epi16(_mm256_maddubs_epi16(aw, _mm256_sign_epi8(x2, w)), ones);
8934            acc1 += hsum(d1) as f32 * s;
8935            acc2 += hsum(d2) as f32 * s;
8936        }
8937        (acc1, acc2)
8938    }
8939}
8940
8941/// One q8 row range via AVX2 (x86 mirror of `q8_range_sdot`).
8942#[cfg(target_arch = "x86_64")]
8943fn q8_range_avx2(
8944    q: &[u8],
8945    row_scale: &[f32],
8946    act: &SplitAct,
8947    cols: usize,
8948    out_addr: SendMut,
8949    start: usize,
8950    end: usize,
8951) {
8952    for o in start..end {
8953        let v = row_dot_avx2(&q[o * cols..(o + 1) * cols], act) * row_scale[o];
8954        // SAFETY: disjoint row ranges per worker.
8955        unsafe { *out_addr.at(o) = v };
8956    }
8957}
8958
8959/// Two-input q8 row range via AVX2 (x86 mirror of `q8_range2_sdot`).
8960#[cfg(target_arch = "x86_64")]
8961#[allow(clippy::too_many_arguments)]
8962fn q8_range2_avx2(
8963    q: &[u8],
8964    row_scale: &[f32],
8965    a1: &SplitAct,
8966    a2: &SplitAct,
8967    cols: usize,
8968    p1: SendMut,
8969    p2: SendMut,
8970    start: usize,
8971    end: usize,
8972) {
8973    for o in start..end {
8974        let row = &q[o * cols..(o + 1) * cols];
8975        // SAFETY: disjoint row ranges per worker.
8976        unsafe {
8977            *p1.at(o) = row_dot_avx2(row, a1) * row_scale[o];
8978            *p2.at(o) = row_dot_avx2(row, a2) * row_scale[o];
8979        }
8980    }
8981}
8982
8983// ───────────────────── A8W8 SDOT path (port of vmfcore, ×1.78 decode) ─────────────────────
8984
8985/// ARMv8.6 i8mm (smmla): 32 int8 MACs per instruction vs sdot's 16 —
8986/// yet MEASURED 2.4× SLOWER than the blocked sdot on Apple silicon
8987/// (108 vs 264 GF/s): the on-the-fly vcombine packing and the two-
8988/// accumulator dependency chain swamp the MAC advantage, and Apple's
8989/// four SIMD pipes already keep sdot fed. OPT-IN (CMF_I8MM=1) for
8990/// field trials on Cortex-A710/X-class parts with two pipes, where the
8991/// balance may differ; a pre-interleaved weight layout (repack infra)
8992/// is the known path if it ever earns its keep.
8993#[cfg(target_arch = "aarch64")]
8994fn i8mm_enabled() -> bool {
8995    static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
8996    *ON.get_or_init(|| {
8997        std::env::var("CMF_I8MM").map(|v| v == "1").unwrap_or(false)
8998            && std::arch::is_aarch64_feature_detected!("i8mm")
8999    })
9000}
9001
9002/// SDOT enabled? Default ON when the CPU has ARMv8.2 dotprod;
9003/// `CMF_SDOT=0` disables (falls back to i8×f32 NEON).
9004/// (On non-ARM release builds only the test tolerance switch calls it.)
9005#[cfg_attr(not(target_arch = "aarch64"), allow(dead_code))]
9006fn sdot_enabled() -> bool {
9007    use std::sync::OnceLock;
9008    static ON: OnceLock<bool> = OnceLock::new();
9009    *ON.get_or_init(|| {
9010        let want = std::env::var("CMF_SDOT").map(|v| v != "0").unwrap_or(true);
9011        if !want {
9012            return false;
9013        }
9014
9015        #[cfg(target_arch = "aarch64")]
9016        {
9017            if std::arch::is_aarch64_feature_detected!("dotprod") {
9018                return true;
9019            }
9020            #[cfg(target_os = "android")]
9021            {
9022                if let Ok(cpuinfo) = std::fs::read_to_string("/proc/cpuinfo") {
9023                    if cpuinfo.lines().any(|l| {
9024                        (l.starts_with("Features") || l.starts_with("features"))
9025                            && l.contains("asimddp")
9026                    }) {
9027                        return true;
9028                    }
9029                }
9030            }
9031            false
9032        }
9033        #[cfg(not(target_arch = "aarch64"))]
9034        {
9035            false
9036        }
9037    })
9038}
9039
9040/// Two-field activation split (≡ vmfcore `q8_split_prep`): outlier
9041/// channels (>8·rms) are computed exactly in f32; the bulk (outliers
9042/// zeroed → clean absmax) goes through int8 SDOT. Computed ONCE per
9043/// matvec, shared by all rows/workers.
9044struct SplitAct {
9045    xq: Vec<i8>,
9046    sx: f32,
9047    outliers: Vec<(usize, f32)>,
9048    /// Σ xq — the VNNI bias-trick correction (`(w+128)·x` sums need
9049    /// `−128·Σx`); one i32 per split, computed once per matvec.
9050    #[cfg_attr(not(target_arch = "x86_64"), allow(dead_code))]
9051    xsum: i32,
9052}
9053
9054thread_local! {
9055    /// Recycled xq buffers: split_act runs for every matvec (~200/token)
9056    /// and its hidden-size allocation was steady-state heap churn.
9057    static XQ_FREE: std::cell::RefCell<Vec<Vec<i8>>> =
9058        const { std::cell::RefCell::new(Vec::new()) };
9059}
9060
9061impl Drop for SplitAct {
9062    fn drop(&mut self) {
9063        let buf = std::mem::take(&mut self.xq);
9064        if buf.capacity() > 0 {
9065            XQ_FREE.with(|f| {
9066                let mut f = f.borrow_mut();
9067                if f.len() < 16 {
9068                    f.push(buf);
9069                }
9070            });
9071        }
9072    }
9073}
9074
9075thread_local! {
9076    /// One scratch row per WORKER, kept for the life of the thread.
9077    ///
9078    /// The kernels take a row of group scales per dispatch, and a fresh
9079    /// `vec![0f32; gpr]` inside the closure is one allocation per worker per
9080    /// dispatch — on the release checkpoint about six thousand a token, a
9081    /// quarter of everything the benchmark counts.
9082    static KROW: std::cell::RefCell<Vec<f32>> = const { std::cell::RefCell::new(Vec::new()) };
9083}
9084
9085/// Borrow `n` floats of the calling worker's scratch. Nothing inside a
9086/// kernel body borrows it again, which is what keeps the RefCell honest.
9087#[inline]
9088fn with_krow<R>(n: usize, f: impl FnOnce(&mut [f32]) -> R) -> R {
9089    KROW.with(|s| {
9090        let mut b = s.borrow_mut();
9091        if b.len() < n {
9092            b.resize(n, 0.0);
9093        }
9094        f(&mut b[..n])
9095    })
9096}
9097
9098fn split_act(x: &[f32]) -> SplitAct {
9099    let n = x.len();
9100    let rms = (x.iter().map(|&v| (v * v) as f64).sum::<f64>() / n.max(1) as f64).sqrt() as f32;
9101    let thr = 8.0 * rms;
9102    // One pass: collect outliers and the bulk absmax (outliers excluded —
9103    // identical to the old zero-then-fold over a copied buffer, minus the
9104    // full-vector copy).
9105    let mut outliers: Vec<(usize, f32)> = Vec::new();
9106    let mut amax = 0f32;
9107    for (j, &v) in x.iter().enumerate() {
9108        let a = v.abs();
9109        if a > thr {
9110            outliers.push((j, v));
9111        } else if a > amax {
9112            amax = a;
9113        }
9114    }
9115    let sx = if amax > 0.0 { amax / 127.0 } else { 1.0 };
9116    let inv = 1.0 / sx;
9117    let mut xq = XQ_FREE.with(|f| f.borrow_mut().pop()).unwrap_or_default();
9118    xq.clear();
9119    xq.reserve(n);
9120    if outliers.is_empty() {
9121        xq.extend(
9122            x.iter()
9123                .map(|&v| (v * inv).round().clamp(-127.0, 127.0) as i8),
9124        );
9125    } else {
9126        // Outlier slots quantize to 0 (their exact term is added later).
9127        xq.extend(x.iter().map(|&v| {
9128            if v.abs() > thr {
9129                0
9130            } else {
9131                (v * inv).round().clamp(-127.0, 127.0) as i8
9132            }
9133        }));
9134    }
9135    let xsum = xq.iter().map(|&v| v as i32).sum();
9136    SplitAct {
9137        xq,
9138        sx,
9139        outliers,
9140        xsum,
9141    }
9142}
9143
9144fn split_act_q8_2f(x: &[f32], col: &[f32]) -> SplitAct {
9145    let n = x.len();
9146    let rms = (x
9147        .iter()
9148        .zip(col)
9149        .map(|(&a, &c)| {
9150            let v = a * c;
9151            (v * v) as f64
9152        })
9153        .sum::<f64>()
9154        / n.max(1) as f64)
9155        .sqrt() as f32;
9156    let thr = 8.0 * rms;
9157
9158    let mut outliers = Vec::new();
9159    let mut amax = 0f32;
9160    for (j, (&a, &c)) in x.iter().zip(col).enumerate() {
9161        let v = a * c;
9162        let s = v.abs();
9163        if s > thr {
9164            outliers.push((j, v));
9165        } else if s > amax {
9166            amax = s;
9167        }
9168    }
9169
9170    let sx = if amax > 0.0 { amax / 127.0 } else { 1.0 };
9171    let inv = 1.0 / sx;
9172    let mut xq = XQ_FREE.with(|f| f.borrow_mut().pop()).unwrap_or_default();
9173    xq.clear();
9174    xq.reserve(n);
9175    if outliers.is_empty() {
9176        xq.extend(
9177            x.iter()
9178                .zip(col)
9179                .map(|(&a, &c)| ((a * c) * inv).round().clamp(-127.0, 127.0) as i8),
9180        );
9181    } else {
9182        xq.extend(x.iter().zip(col).map(|(&a, &c)| {
9183            let v = a * c;
9184            if v.abs() > thr {
9185                0
9186            } else {
9187                (v * inv).round().clamp(-127.0, 127.0) as i8
9188            }
9189        }));
9190    }
9191    let xsum = xq.iter().map(|&v| v as i32).sum();
9192    SplitAct {
9193        xq,
9194        sx,
9195        outliers,
9196        xsum,
9197    }
9198}
9199
9200/// int8(weight)·int8(activation) → i32 via `sdot` (inline asm — the
9201/// vdotq intrinsic is unstable; port of vmfcore `dot_i8_sdot`).
9202#[cfg(target_arch = "aarch64")]
9203#[target_feature(enable = "neon,dotprod")]
9204unsafe fn dot_i8_sdot(w: &[u8], xq: &[i8]) -> i32 {
9205    // SAFETY: callers uphold slice-length contracts (see call sites).
9206    unsafe {
9207        use core::arch::aarch64::*;
9208        use core::arch::asm;
9209        let wp = w.as_ptr() as *const i8;
9210        let n = w.len();
9211        let (mut a0, mut a1, mut a2, mut a3) = (
9212            vdupq_n_s32(0),
9213            vdupq_n_s32(0),
9214            vdupq_n_s32(0),
9215            vdupq_n_s32(0),
9216        );
9217        let mut i = 0;
9218        while i + 64 <= n {
9219            let (w0, x0) = (vld1q_s8(wp.add(i)), vld1q_s8(xq.as_ptr().add(i)));
9220            let (w1, x1) = (vld1q_s8(wp.add(i + 16)), vld1q_s8(xq.as_ptr().add(i + 16)));
9221            let (w2, x2) = (vld1q_s8(wp.add(i + 32)), vld1q_s8(xq.as_ptr().add(i + 32)));
9222            let (w3, x3) = (vld1q_s8(wp.add(i + 48)), vld1q_s8(xq.as_ptr().add(i + 48)));
9223            asm!(
9224                "sdot {a0:v}.4s, {w0:v}.16b, {x0:v}.16b",
9225                "sdot {a1:v}.4s, {w1:v}.16b, {x1:v}.16b",
9226                "sdot {a2:v}.4s, {w2:v}.16b, {x2:v}.16b",
9227                "sdot {a3:v}.4s, {w3:v}.16b, {x3:v}.16b",
9228                a0 = inout(vreg) a0, a1 = inout(vreg) a1, a2 = inout(vreg) a2, a3 = inout(vreg) a3,
9229                w0 = in(vreg) w0, x0 = in(vreg) x0, w1 = in(vreg) w1, x1 = in(vreg) x1,
9230                w2 = in(vreg) w2, x2 = in(vreg) x2, w3 = in(vreg) w3, x3 = in(vreg) x3,
9231                options(pure, nomem, nostack),
9232            );
9233            i += 64;
9234        }
9235        while i + 16 <= n {
9236            let (wv, xv) = (vld1q_s8(wp.add(i)), vld1q_s8(xq.as_ptr().add(i)));
9237            asm!("sdot {a:v}.4s, {w:v}.16b, {x:v}.16b",
9238                 a = inout(vreg) a0, w = in(vreg) wv, x = in(vreg) xv, options(pure, nomem, nostack));
9239            i += 16;
9240        }
9241        let mut s = vaddvq_s32(vaddq_s32(vaddq_s32(a0, a1), vaddq_s32(a2, a3)));
9242        while i < n {
9243            s += (*wp.add(i)) as i32 * xq[i] as i32;
9244            i += 1;
9245        }
9246        s
9247    }
9248}
9249
9250/// Row-blocked SDOT: 4 output rows per pass — the activation chunk is
9251/// loaded once and reused, 4 independent accumulators hide sdot latency
9252/// (port of vmfcore `dot_i8_sdot_4rows`).
9253#[cfg(target_arch = "aarch64")]
9254#[target_feature(enable = "neon,dotprod")]
9255unsafe fn dot_i8_sdot_4rows(w0: &[u8], w1: &[u8], w2: &[u8], w3: &[u8], xq: &[i8]) -> [i32; 4] {
9256    // SAFETY: callers uphold slice-length contracts (see call sites).
9257    unsafe {
9258        use core::arch::aarch64::*;
9259        use core::arch::asm;
9260        let n = xq.len();
9261        let px = xq.as_ptr();
9262        let (p0, p1, p2, p3) = (
9263            w0.as_ptr() as *const i8,
9264            w1.as_ptr() as *const i8,
9265            w2.as_ptr() as *const i8,
9266            w3.as_ptr() as *const i8,
9267        );
9268        let (mut a0, mut a1, mut a2, mut a3) = (
9269            vdupq_n_s32(0),
9270            vdupq_n_s32(0),
9271            vdupq_n_s32(0),
9272            vdupq_n_s32(0),
9273        );
9274        let mut i = 0;
9275        while i + 16 <= n {
9276            let x = vld1q_s8(px.add(i));
9277            let v0 = vld1q_s8(p0.add(i));
9278            let v1 = vld1q_s8(p1.add(i));
9279            let v2 = vld1q_s8(p2.add(i));
9280            let v3 = vld1q_s8(p3.add(i));
9281            asm!(
9282                "sdot {a0:v}.4s, {v0:v}.16b, {x:v}.16b",
9283                "sdot {a1:v}.4s, {v1:v}.16b, {x:v}.16b",
9284                "sdot {a2:v}.4s, {v2:v}.16b, {x:v}.16b",
9285                "sdot {a3:v}.4s, {v3:v}.16b, {x:v}.16b",
9286                a0 = inout(vreg) a0, a1 = inout(vreg) a1, a2 = inout(vreg) a2, a3 = inout(vreg) a3,
9287                v0 = in(vreg) v0, v1 = in(vreg) v1, v2 = in(vreg) v2, v3 = in(vreg) v3, x = in(vreg) x,
9288                options(pure, nomem, nostack),
9289            );
9290            i += 16;
9291        }
9292        let mut r = [
9293            vaddvq_s32(a0),
9294            vaddvq_s32(a1),
9295            vaddvq_s32(a2),
9296            vaddvq_s32(a3),
9297        ];
9298        while i < n {
9299            let xi = *px.add(i) as i32;
9300            r[0] += (*p0.add(i)) as i32 * xi;
9301            r[1] += (*p1.add(i)) as i32 * xi;
9302            r[2] += (*p2.add(i)) as i32 * xi;
9303            r[3] += (*p3.add(i)) as i32 * xi;
9304            i += 1;
9305        }
9306        r
9307    }
9308}
9309
9310/// 4 interleaved rows in one pass: the repacked group is [r0[c], r1[c],
9311/// r2[c], r3[c]] per 16-byte chunk, so each iteration reads ONE 64-byte
9312/// line plus the shared activation chunk — a single sequential weight
9313/// stream per worker. Per-row accumulation is the same one-accumulator
9314/// scheme as `dot_i8_sdot_4rows`; integer sums are exact, so outputs
9315/// are bit-identical to the mmap-layout kernel.
9316#[cfg(target_arch = "aarch64")]
9317#[target_feature(enable = "neon,dotprod")]
9318unsafe fn dot_i8_sdot_4rows_il(g: &[u8], xq: &[i8]) -> [i32; 4] {
9319    // SAFETY: callers uphold slice-length contracts (g.len() == 4·n,
9320    // n % 16 == 0 — guaranteed by the repack gate).
9321    unsafe {
9322        use core::arch::aarch64::*;
9323        use core::arch::asm;
9324        let n = xq.len();
9325        let px = xq.as_ptr();
9326        let pg = g.as_ptr() as *const i8;
9327        let (mut a0, mut a1, mut a2, mut a3) = (
9328            vdupq_n_s32(0),
9329            vdupq_n_s32(0),
9330            vdupq_n_s32(0),
9331            vdupq_n_s32(0),
9332        );
9333        let mut i = 0;
9334        while i + 16 <= n {
9335            let x = vld1q_s8(px.add(i));
9336            let base = pg.add(4 * i);
9337            let v0 = vld1q_s8(base);
9338            let v1 = vld1q_s8(base.add(16));
9339            let v2 = vld1q_s8(base.add(32));
9340            let v3 = vld1q_s8(base.add(48));
9341            asm!(
9342                "sdot {a0:v}.4s, {v0:v}.16b, {x:v}.16b",
9343                "sdot {a1:v}.4s, {v1:v}.16b, {x:v}.16b",
9344                "sdot {a2:v}.4s, {v2:v}.16b, {x:v}.16b",
9345                "sdot {a3:v}.4s, {v3:v}.16b, {x:v}.16b",
9346                a0 = inout(vreg) a0, a1 = inout(vreg) a1, a2 = inout(vreg) a2, a3 = inout(vreg) a3,
9347                v0 = in(vreg) v0, v1 = in(vreg) v1, v2 = in(vreg) v2, v3 = in(vreg) v3, x = in(vreg) x,
9348                options(pure, nomem, nostack),
9349            );
9350            i += 16;
9351        }
9352        [
9353            vaddvq_s32(a0),
9354            vaddvq_s32(a1),
9355            vaddvq_s32(a2),
9356            vaddvq_s32(a3),
9357        ]
9358    }
9359}
9360
9361/// One q8 row range via SDOT (4-row blocks + tail) — the body of
9362/// `qmatvec`'s hot loop, extracted so multi-matrix jobs can drive the
9363/// SAME kernel for several tensors under one pool dispatch. `rep` — the
9364/// load-time interleaved repack (empty = mmap layout only); rows outside
9365/// full 4-row groups always come from the mmap layout.
9366#[cfg(target_arch = "aarch64")]
9367fn q8_range_sdot(
9368    q: &[u8],
9369    rep: &[u8],
9370    row_scale: &[f32],
9371    act: &SplitAct,
9372    cols: usize,
9373    out_addr: SendMut,
9374    start: usize,
9375    end: usize,
9376) {
9377    let mut o = start;
9378    // Leading rows to the group boundary (repack path only): the pool
9379    // splits row ranges arbitrarily, groups are absolute.
9380    if !rep.is_empty() {
9381        while o < end && o % 4 != 0 {
9382            let v = row_dot_sdot(&q[o * cols..(o + 1) * cols], act) * row_scale[o];
9383            unsafe { *out_addr.at(o) = v };
9384            o += 1;
9385        }
9386    }
9387    while o + 4 <= end {
9388        let r = if rep.is_empty() {
9389            unsafe {
9390                dot_i8_sdot_4rows(
9391                    &q[o * cols..(o + 1) * cols],
9392                    &q[(o + 1) * cols..(o + 2) * cols],
9393                    &q[(o + 2) * cols..(o + 3) * cols],
9394                    &q[(o + 3) * cols..(o + 4) * cols],
9395                    &act.xq,
9396                )
9397            }
9398        } else {
9399            unsafe { dot_i8_sdot_4rows_il(&rep[o * cols..(o + 4) * cols], &act.xq) }
9400        };
9401        for k in 0..4 {
9402            let mut acc = r[k] as f32 * act.sx;
9403            for &(j, xv) in &act.outliers {
9404                acc += (q[(o + k) * cols + j] as i8) as f32 * xv;
9405            }
9406            // SAFETY: disjoint row ranges per worker.
9407            unsafe { *out_addr.at(o + k) = acc * row_scale[o + k] };
9408        }
9409        o += 4;
9410    }
9411    while o < end {
9412        let v = row_dot_sdot(&q[o * cols..(o + 1) * cols], act) * row_scale[o];
9413        unsafe { *out_addr.at(o) = v };
9414        o += 1;
9415    }
9416}
9417
9418/// Two-input q8 row range via SDOT — `qmatvec2`'s hot loop, extracted
9419/// for the fused pair multi-matrix job (`matvec2_many`).
9420#[cfg(target_arch = "aarch64")]
9421#[allow(clippy::too_many_arguments)]
9422fn q8_range2_sdot(
9423    q: &[u8],
9424    row_scale: &[f32],
9425    a1: &SplitAct,
9426    a2: &SplitAct,
9427    cols: usize,
9428    p1: SendMut,
9429    p2: SendMut,
9430    start: usize,
9431    end: usize,
9432) {
9433    for o in start..end {
9434        let row = &q[o * cols..(o + 1) * cols];
9435        // SAFETY: disjoint row ranges per worker.
9436        unsafe {
9437            *p1.at(o) = row_dot_sdot(row, a1) * row_scale[o];
9438            *p2.at(o) = row_dot_sdot(row, a2) * row_scale[o];
9439        }
9440    }
9441}
9442
9443/// Two-input q8 row range, f32 kernel (non-SDOT) — same extraction.
9444#[allow(clippy::too_many_arguments)]
9445fn q8_range2_f32(
9446    q: &[u8],
9447    row_scale: &[f32],
9448    x1: &[f32],
9449    x2: &[f32],
9450    cols: usize,
9451    p1: SendMut,
9452    p2: SendMut,
9453    start: usize,
9454    end: usize,
9455) {
9456    for o in start..end {
9457        let row = &q[o * cols..(o + 1) * cols];
9458        // SAFETY: disjoint row ranges per worker.
9459        unsafe {
9460            *p1.at(o) = dot_i8_f32(row, x1) * row_scale[o];
9461            *p2.at(o) = dot_i8_f32(row, x2) * row_scale[o];
9462        }
9463    }
9464}
9465
9466/// Scalar/NEON-f32 q8 row range (non-SDOT platforms) — same extraction.
9467fn q8_range_f32(
9468    q: &[u8],
9469    row_scale: &[f32],
9470    xs: &[f32],
9471    cols: usize,
9472    out_addr: SendMut,
9473    start: usize,
9474    end: usize,
9475) {
9476    for o in start..end {
9477        let v = dot_i8_f32(&q[o * cols..(o + 1) * cols], xs) * row_scale[o];
9478        // SAFETY: disjoint row ranges per worker.
9479        unsafe { *out_addr.at(o) = v };
9480    }
9481}
9482
9483/// One q8 row against a split activation, portable: the per-arch fast
9484/// dots where they exist, the exact scalar loop elsewhere. The scalar
9485/// arm is also the test oracle for both fast arms.
9486#[inline]
9487fn q8_row_dot(row: &[u8], act: &SplitAct) -> f32 {
9488    #[cfg(target_arch = "aarch64")]
9489    return row_dot_sdot(row, act);
9490    #[cfg(target_arch = "x86_64")]
9491    return row_dot_avx2(row, act);
9492    #[allow(unreachable_code)]
9493    q8_row_dot_scalar(row, act)
9494}
9495
9496#[allow(dead_code)]
9497fn q8_row_dot_scalar(row: &[u8], act: &SplitAct) -> f32 {
9498    let mut acc = 0i32;
9499    for (k, &b) in row.iter().enumerate() {
9500        acc += (b as i8) as i32 * act.xq[k] as i32;
9501    }
9502    let mut acc = acc as f32 * act.sx;
9503    for &(j, xv) in &act.outliers {
9504        acc += (row[j] as i8) as f32 * xv;
9505    }
9506    acc
9507}
9508
9509/// SDOT row dot with exact outlier correction:
9510/// `dot = sdot(w, xq)·sx + Σ_outl w[j]·x[j]` (then × row_scale by caller).
9511#[cfg(target_arch = "aarch64")]
9512#[inline]
9513fn row_dot_sdot(row: &[u8], act: &SplitAct) -> f32 {
9514    let mut acc = unsafe { dot_i8_sdot(row, &act.xq) } as f32 * act.sx;
9515    for &(j, xv) in &act.outliers {
9516        acc += (row[j] as i8) as f32 * xv;
9517    }
9518    acc
9519}
9520
9521/// One q4 row via SDOT: each 32-group's nibbles unpack to centered i8
9522/// (nib−8 ∈ [−8,7]), int8×int8 `sdot` against the pre-quantized
9523/// activation group, × the group's f16 scale. Returns Σ_g dot_g·s_g;
9524/// the caller multiplies by the activation scale and adds the exact
9525/// outlier terms (port of vmfcore `dot_q4_block_sdot`, +23% measured).
9526/// Nibble order matches the writer: element 2k = low nibble, 2k+1 = high
9527/// → zip(lo,hi) restores flat order.
9528#[cfg(target_arch = "aarch64")]
9529#[target_feature(enable = "neon,dotprod")]
9530unsafe fn dot_q4_row_sdot(packed: &[u8], scales: &[u8], g0: usize, gpr: usize, xq: &[i8]) -> f32 {
9531    // SAFETY: callers uphold slice-length contracts (16 packed bytes and
9532    // 2 scale bytes per group; xq.len() == gpr·GROUP_SIZE).
9533    unsafe {
9534        use core::arch::aarch64::*;
9535        use core::arch::asm;
9536        let lomask = vdupq_n_u8(0x0F);
9537        let eight = vdupq_n_s8(8);
9538        let mut acc = 0f32;
9539        for gi in 0..gpr {
9540            let g = g0 + gi;
9541            let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
9542            let b = vld1q_u8(packed.as_ptr().add(g * 16));
9543            let lo = vandq_u8(b, lomask);
9544            let hi = vshrq_n_u8::<4>(b);
9545            let e0 = vsubq_s8(vreinterpretq_s8_u8(vzip1q_u8(lo, hi)), eight);
9546            let e1 = vsubq_s8(vreinterpretq_s8_u8(vzip2q_u8(lo, hi)), eight);
9547            let x0 = vld1q_s8(xq.as_ptr().add(gi * GROUP_SIZE));
9548            let x1 = vld1q_s8(xq.as_ptr().add(gi * GROUP_SIZE + 16));
9549            let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
9550            asm!(
9551                "sdot {a0:v}.4s, {e0:v}.16b, {x0:v}.16b",
9552                "sdot {a1:v}.4s, {e1:v}.16b, {x1:v}.16b",
9553                a0 = inout(vreg) a0, a1 = inout(vreg) a1,
9554                e0 = in(vreg) e0, x0 = in(vreg) x0, e1 = in(vreg) e1, x1 = in(vreg) x1,
9555                options(pure, nomem, nostack),
9556            );
9557            acc += vaddvq_s32(vaddq_s32(a0, a1)) as f32 * s;
9558        }
9559        acc
9560    }
9561}
9562
9563/// Two-activation q4 row via SDOT: the nibble unpack (the expensive
9564/// part) happens ONCE per group; both pre-quantized activations are
9565/// dotted against the same centered i8 registers. Per-lane math matches
9566/// `dot_q4_row_sdot` exactly.
9567#[cfg(target_arch = "aarch64")]
9568#[target_feature(enable = "neon,dotprod")]
9569unsafe fn dot_q4_row_sdot2(
9570    packed: &[u8],
9571    scales: &[u8],
9572    g0: usize,
9573    gpr: usize,
9574    xq1: &[i8],
9575    xq2: &[i8],
9576) -> (f32, f32) {
9577    // SAFETY: callers uphold slice-length contracts (16 packed bytes and
9578    // 2 scale bytes per group; xq*.len() == gpr·GROUP_SIZE).
9579    unsafe {
9580        use core::arch::aarch64::*;
9581        use core::arch::asm;
9582        let lomask = vdupq_n_u8(0x0F);
9583        let eight = vdupq_n_s8(8);
9584        let (mut acc1, mut acc2) = (0f32, 0f32);
9585        for gi in 0..gpr {
9586            let g = g0 + gi;
9587            let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
9588            let b = vld1q_u8(packed.as_ptr().add(g * 16));
9589            let lo = vandq_u8(b, lomask);
9590            let hi = vshrq_n_u8::<4>(b);
9591            let e0 = vsubq_s8(vreinterpretq_s8_u8(vzip1q_u8(lo, hi)), eight);
9592            let e1 = vsubq_s8(vreinterpretq_s8_u8(vzip2q_u8(lo, hi)), eight);
9593            let x10 = vld1q_s8(xq1.as_ptr().add(gi * GROUP_SIZE));
9594            let x11 = vld1q_s8(xq1.as_ptr().add(gi * GROUP_SIZE + 16));
9595            let x20 = vld1q_s8(xq2.as_ptr().add(gi * GROUP_SIZE));
9596            let x21 = vld1q_s8(xq2.as_ptr().add(gi * GROUP_SIZE + 16));
9597            let (mut a0, mut a1, mut b0, mut b1) = (
9598                vdupq_n_s32(0),
9599                vdupq_n_s32(0),
9600                vdupq_n_s32(0),
9601                vdupq_n_s32(0),
9602            );
9603            asm!(
9604                "sdot {a0:v}.4s, {e0:v}.16b, {x10:v}.16b",
9605                "sdot {a1:v}.4s, {e1:v}.16b, {x11:v}.16b",
9606                "sdot {b0:v}.4s, {e0:v}.16b, {x20:v}.16b",
9607                "sdot {b1:v}.4s, {e1:v}.16b, {x21:v}.16b",
9608                a0 = inout(vreg) a0, a1 = inout(vreg) a1,
9609                b0 = inout(vreg) b0, b1 = inout(vreg) b1,
9610                e0 = in(vreg) e0, e1 = in(vreg) e1,
9611                x10 = in(vreg) x10, x11 = in(vreg) x11,
9612                x20 = in(vreg) x20, x21 = in(vreg) x21,
9613                options(pure, nomem, nostack),
9614            );
9615            acc1 += vaddvq_s32(vaddq_s32(a0, a1)) as f32 * s;
9616            acc2 += vaddvq_s32(vaddq_s32(b0, b1)) as f32 * s;
9617        }
9618        (acc1, acc2)
9619    }
9620}
9621
9622// ───────────────────── fused int8 kernels ─────────────────────
9623
9624/// `acc += w · row` where the row is centered i8 — NEON widen+fma on
9625/// aarch64, scalar elsewhere. The KV-cache q8 value path rides on this.
9626#[inline]
9627pub(crate) fn axpy_i8_f32(acc: &mut [f32], row: &[i8], w: f32) {
9628    #[cfg(target_arch = "aarch64")]
9629    unsafe {
9630        return axpy_i8_f32_neon(acc, row, w);
9631    }
9632    #[cfg(target_arch = "x86_64")]
9633    if avx2_enabled() {
9634        return unsafe { axpy_i8_f32_avx2(acc, row, w) };
9635    }
9636    #[allow(unreachable_code)]
9637    {
9638        for (a, &b) in acc.iter_mut().zip(row) {
9639            *a += w * b as f32;
9640        }
9641    }
9642}
9643
9644/// i8→f32 axpy via AVX2/FMA (x86 mirror of `axpy_i8_f32_neon`).
9645#[cfg(target_arch = "x86_64")]
9646#[target_feature(enable = "avx2,fma")]
9647unsafe fn axpy_i8_f32_avx2(acc: &mut [f32], row: &[i8], w: f32) {
9648    // SAFETY: callers uphold slice-length contracts (see call sites).
9649    unsafe {
9650        use core::arch::x86_64::*;
9651        let n = acc.len().min(row.len());
9652        let ap = acc.as_mut_ptr();
9653        let rp = row.as_ptr();
9654        let wv = _mm256_set1_ps(w);
9655        let mut j = 0usize;
9656        while j + 16 <= n {
9657            let rb = _mm_loadu_si128(rp.add(j) as *const __m128i);
9658            let lo = _mm256_cvtepi32_ps(_mm256_cvtepi8_epi32(rb));
9659            let hi = _mm256_cvtepi32_ps(_mm256_cvtepi8_epi32(_mm_srli_si128::<8>(rb)));
9660            let v0 = _mm256_fmadd_ps(wv, lo, _mm256_loadu_ps(ap.add(j)));
9661            let v1 = _mm256_fmadd_ps(wv, hi, _mm256_loadu_ps(ap.add(j + 8)));
9662            _mm256_storeu_ps(ap.add(j), v0);
9663            _mm256_storeu_ps(ap.add(j + 8), v1);
9664            j += 16;
9665        }
9666        while j < n {
9667            *ap.add(j) += w * (*rp.add(j)) as f32;
9668            j += 1;
9669        }
9670    }
9671}
9672
9673#[cfg(target_arch = "aarch64")]
9674#[target_feature(enable = "neon")]
9675unsafe fn axpy_i8_f32_neon(acc: &mut [f32], row: &[i8], w: f32) {
9676    // SAFETY: callers uphold slice-length contracts (see call sites).
9677    unsafe {
9678        use core::arch::aarch64::*;
9679        let n = acc.len().min(row.len());
9680        let ap = acc.as_mut_ptr();
9681        let rp = row.as_ptr();
9682        let wv = vdupq_n_f32(w);
9683        let mut j = 0usize;
9684        while j + 16 <= n {
9685            let rb = vld1q_s8(rp.add(j));
9686            let lo = vmovl_s8(vget_low_s8(rb));
9687            let hi = vmovl_s8(vget_high_s8(rb));
9688            for (off, half) in [(0, lo), (8, hi)] {
9689                let f0 = vcvtq_f32_s32(vmovl_s16(vget_low_s16(half)));
9690                let f1 = vcvtq_f32_s32(vmovl_s16(vget_high_s16(half)));
9691                let o = j + off;
9692                vst1q_f32(ap.add(o), vfmaq_f32(vld1q_f32(ap.add(o)), wv, f0));
9693                vst1q_f32(ap.add(o + 4), vfmaq_f32(vld1q_f32(ap.add(o + 4)), wv, f1));
9694            }
9695            j += 16;
9696        }
9697        while j < n {
9698            *ap.add(j) += w * (*rp.add(j)) as f32;
9699            j += 1;
9700        }
9701    }
9702}
9703
9704/// i8 row · f32 x. NEON on aarch64 (ported from vmfcore `dot_i8_f32_neon`,
9705/// ≈9× scalar), scalar elsewhere.
9706#[inline]
9707pub(crate) fn dot_i8_f32(w: &[u8], x: &[f32]) -> f32 {
9708    #[cfg(target_arch = "aarch64")]
9709    unsafe {
9710        return dot_i8_f32_neon(w, x);
9711    }
9712    #[cfg(target_arch = "x86_64")]
9713    if avx2_enabled() {
9714        return unsafe { dot_i8_f32_avx2(w, x) };
9715    }
9716    #[allow(unreachable_code)]
9717    {
9718        let mut sum = 0.0f32;
9719        for (j, &b) in w.iter().enumerate() {
9720            sum += (b as i8) as f32 * x[j];
9721        }
9722        sum
9723    }
9724}
9725
9726/// i8 row · (x ⊙ col_field) — the q8_2f row dot with the θ col-field
9727/// folded into the product (no prescaled copy of x). NEON on aarch64,
9728/// scalar elsewhere. Used by the active-neuron path `row_dot`.
9729#[inline]
9730fn dot_i8_col_f32(w: &[u8], x: &[f32], col: &[f32]) -> f32 {
9731    #[cfg(target_arch = "aarch64")]
9732    unsafe {
9733        return dot_i8_col_f32_neon(w, x, col);
9734    }
9735    #[allow(unreachable_code)]
9736    {
9737        let mut sum = 0.0f32;
9738        for (j, &b) in w.iter().enumerate() {
9739            sum += (b as i8) as f32 * x[j] * col[j];
9740        }
9741        sum
9742    }
9743}
9744
9745#[cfg(target_arch = "aarch64")]
9746#[target_feature(enable = "neon")]
9747unsafe fn dot_i8_col_f32_neon(w: &[u8], x: &[f32], col: &[f32]) -> f32 {
9748    // SAFETY: callers uphold slice-length contracts (see call sites).
9749    unsafe {
9750        use core::arch::aarch64::*;
9751        let n = x.len();
9752        let wp = w.as_ptr() as *const i8;
9753        let xp = x.as_ptr();
9754        let cp = col.as_ptr();
9755        let (mut a0, mut a1, mut a2, mut a3) = (
9756            vdupq_n_f32(0.0),
9757            vdupq_n_f32(0.0),
9758            vdupq_n_f32(0.0),
9759            vdupq_n_f32(0.0),
9760        );
9761        let mut j = 0usize;
9762        while j + 16 <= n {
9763            let wb = vld1q_s8(wp.add(j));
9764            let lo = vmovl_s8(vget_low_s8(wb));
9765            let hi = vmovl_s8(vget_high_s8(wb));
9766            let w0 = vcvtq_f32_s32(vmovl_s16(vget_low_s16(lo)));
9767            let w1 = vcvtq_f32_s32(vmovl_s16(vget_high_s16(lo)));
9768            let w2 = vcvtq_f32_s32(vmovl_s16(vget_low_s16(hi)));
9769            let w3 = vcvtq_f32_s32(vmovl_s16(vget_high_s16(hi)));
9770            a0 = vfmaq_f32(
9771                a0,
9772                w0,
9773                vmulq_f32(vld1q_f32(xp.add(j)), vld1q_f32(cp.add(j))),
9774            );
9775            a1 = vfmaq_f32(
9776                a1,
9777                w1,
9778                vmulq_f32(vld1q_f32(xp.add(j + 4)), vld1q_f32(cp.add(j + 4))),
9779            );
9780            a2 = vfmaq_f32(
9781                a2,
9782                w2,
9783                vmulq_f32(vld1q_f32(xp.add(j + 8)), vld1q_f32(cp.add(j + 8))),
9784            );
9785            a3 = vfmaq_f32(
9786                a3,
9787                w3,
9788                vmulq_f32(vld1q_f32(xp.add(j + 12)), vld1q_f32(cp.add(j + 12))),
9789            );
9790            j += 16;
9791        }
9792        let mut sum = vaddvq_f32(vaddq_f32(vaddq_f32(a0, a1), vaddq_f32(a2, a3)));
9793        while j < n {
9794            sum += (*wp.add(j)) as f32 * *xp.add(j) * *cp.add(j);
9795            j += 1;
9796        }
9797        sum
9798    }
9799}
9800
9801#[cfg(target_arch = "aarch64")]
9802#[target_feature(enable = "neon")]
9803unsafe fn dot_i8_f32_neon(w: &[u8], x: &[f32]) -> f32 {
9804    // SAFETY: callers uphold slice-length contracts (see call sites).
9805    unsafe {
9806        use core::arch::aarch64::*;
9807        let n = x.len();
9808        let wp = w.as_ptr() as *const i8;
9809        let xp = x.as_ptr();
9810        let (mut a0, mut a1, mut a2, mut a3) = (
9811            vdupq_n_f32(0.0),
9812            vdupq_n_f32(0.0),
9813            vdupq_n_f32(0.0),
9814            vdupq_n_f32(0.0),
9815        );
9816        let mut j = 0usize;
9817        while j + 16 <= n {
9818            let wb = vld1q_s8(wp.add(j));
9819            let lo = vmovl_s8(vget_low_s8(wb));
9820            let hi = vmovl_s8(vget_high_s8(wb));
9821            let w0 = vcvtq_f32_s32(vmovl_s16(vget_low_s16(lo)));
9822            let w1 = vcvtq_f32_s32(vmovl_s16(vget_high_s16(lo)));
9823            let w2 = vcvtq_f32_s32(vmovl_s16(vget_low_s16(hi)));
9824            let w3 = vcvtq_f32_s32(vmovl_s16(vget_high_s16(hi)));
9825            a0 = vfmaq_f32(a0, w0, vld1q_f32(xp.add(j)));
9826            a1 = vfmaq_f32(a1, w1, vld1q_f32(xp.add(j + 4)));
9827            a2 = vfmaq_f32(a2, w2, vld1q_f32(xp.add(j + 8)));
9828            a3 = vfmaq_f32(a3, w3, vld1q_f32(xp.add(j + 12)));
9829            j += 16;
9830        }
9831        let mut sum = vaddvq_f32(vaddq_f32(vaddq_f32(a0, a1), vaddq_f32(a2, a3)));
9832        while j < n {
9833            sum += (*wp.add(j)) as f32 * *xp.add(j);
9834            j += 1;
9835        }
9836        sum
9837    }
9838}
9839
9840#[allow(clippy::too_many_arguments)]
9841fn qmatvec(
9842    q: &[u8],
9843    rep: &[u8],
9844    row_scale: &[f32],
9845    x: &[f32],
9846    col_field: &[f32],
9847    dtype: TensorDtype,
9848    rows: usize,
9849    cols: usize,
9850    out: &mut [f32],
9851    pool: Option<&Pool>,
9852) {
9853    debug_assert_eq!(out.len(), rows);
9854    #[cfg(not(target_arch = "aarch64"))]
9855    let _ = rep;
9856
9857    #[cfg(target_arch = "aarch64")]
9858    if sdot_enabled() {
9859        let act = if dtype == TensorDtype::Q8_2f {
9860            split_act_q8_2f(x, col_field)
9861        } else {
9862            split_act(x)
9863        };
9864        let out_addr = SendMut(out.as_mut_ptr());
9865        let run_range = |start: usize, end: usize| {
9866            q8_range_sdot(q, rep, row_scale, &act, cols, out_addr, start, end)
9867        };
9868        match pool {
9869            Some(pool) if rows >= 256 => pool.run_rows(rows, &run_range),
9870            _ => run_range(0, rows),
9871        }
9872        return;
9873    }
9874    // x86 A8W8 via AVX2 maddubs — same quantized-activation contract as
9875    // the SDOT path (CMF_AVX2=0 keeps the exact i8×f32 loop).
9876    #[cfg(target_arch = "x86_64")]
9877    if avx2_a8w8_enabled() {
9878        let act = if dtype == TensorDtype::Q8_2f {
9879            split_act_q8_2f(x, col_field)
9880        } else {
9881            split_act(x)
9882        };
9883        let out_addr = SendMut(out.as_mut_ptr());
9884        let run_range = |start: usize, end: usize| {
9885            q8_range_avx2(q, row_scale, &act, cols, out_addr, start, end)
9886        };
9887        match pool {
9888            Some(pool) if rows >= 256 => pool.run_rows(rows, &run_range),
9889            _ => run_range(0, rows),
9890        }
9891        return;
9892    }
9893
9894    prescale_with(x, col_field, dtype, 1, |xs| {
9895        let out_addr = SendMut(out.as_mut_ptr());
9896        let run_range = move |start: usize, end: usize| {
9897            for o in start..end {
9898                let v = dot_i8_f32(&q[o * cols..(o + 1) * cols], xs) * row_scale[o];
9899                // SAFETY: disjoint row ranges per worker.
9900                unsafe { *out_addr.at(o) = v };
9901            }
9902        };
9903        match pool {
9904            Some(pool) if rows >= 256 => pool.run_rows(rows, &run_range),
9905            _ => run_range(0, rows),
9906        }
9907    });
9908}
9909
9910#[allow(clippy::too_many_arguments)]
9911fn qmatvec2(
9912    q: &[u8],
9913    row_scale: &[f32],
9914    x1: &[f32],
9915    x2: &[f32],
9916    col_field: &[f32],
9917    dtype: TensorDtype,
9918    rows: usize,
9919    cols: usize,
9920    o1: &mut [f32],
9921    o2: &mut [f32],
9922    pool: Option<&Pool>,
9923) {
9924    #[cfg(target_arch = "aarch64")]
9925    if sdot_enabled() {
9926        let a1s = if dtype == TensorDtype::Q8_2f {
9927            split_act_q8_2f(x1, col_field)
9928        } else {
9929            split_act(x1)
9930        };
9931        let a2s = if dtype == TensorDtype::Q8_2f {
9932            split_act_q8_2f(x2, col_field)
9933        } else {
9934            split_act(x2)
9935        };
9936        let p1 = SendMut(o1.as_mut_ptr());
9937        let p2 = SendMut(o2.as_mut_ptr());
9938        let run_range = |start: usize, end: usize| {
9939            q8_range2_sdot(q, row_scale, &a1s, &a2s, cols, p1, p2, start, end)
9940        };
9941        match pool {
9942            Some(pool) if rows >= 256 => pool.run_rows(rows, &run_range),
9943            _ => run_range(0, rows),
9944        }
9945        return;
9946    }
9947    #[cfg(target_arch = "x86_64")]
9948    if avx2_a8w8_enabled() {
9949        let a1s = if dtype == TensorDtype::Q8_2f {
9950            split_act_q8_2f(x1, col_field)
9951        } else {
9952            split_act(x1)
9953        };
9954        let a2s = if dtype == TensorDtype::Q8_2f {
9955            split_act_q8_2f(x2, col_field)
9956        } else {
9957            split_act(x2)
9958        };
9959        let p1 = SendMut(o1.as_mut_ptr());
9960        let p2 = SendMut(o2.as_mut_ptr());
9961        let run_range = |start: usize, end: usize| {
9962            q8_range2_avx2(q, row_scale, &a1s, &a2s, cols, p1, p2, start, end)
9963        };
9964        match pool {
9965            Some(pool) if rows >= 256 => pool.run_rows(rows, &run_range),
9966            _ => run_range(0, rows),
9967        }
9968        return;
9969    }
9970
9971    prescale_with(x1, col_field, dtype, 1, |x1s| {
9972        prescale_with(x2, col_field, dtype, 2, |x2s| {
9973            let p1 = SendMut(o1.as_mut_ptr());
9974            let p2 = SendMut(o2.as_mut_ptr());
9975            let run_range = move |start: usize, end: usize| {
9976                for o in start..end {
9977                    let row = &q[o * cols..(o + 1) * cols];
9978                    let s1 = dot_i8_f32(row, x1s) * row_scale[o];
9979                    let s2 = dot_i8_f32(row, x2s) * row_scale[o];
9980                    // SAFETY: disjoint row ranges per worker.
9981                    unsafe {
9982                        *p1.at(o) = s1;
9983                        *p2.at(o) = s2;
9984                    }
9985                }
9986            };
9987            match pool {
9988                Some(pool) if rows >= 256 => pool.run_rows(rows, &run_range),
9989                _ => run_range(0, rows),
9990            }
9991        });
9992    });
9993}
9994
9995#[derive(Clone, Copy)]
9996struct SendMut(*mut f32);
9997unsafe impl Send for SendMut {}
9998unsafe impl Sync for SendMut {}
9999
10000impl SendMut {
10001    #[inline]
10002    fn at(self, i: usize) -> *mut f32 {
10003        unsafe { self.0.add(i) }
10004    }
10005}
10006
10007#[cfg(test)]
10008mod tests {
10009    use super::*;
10010
10011    #[test]
10012    fn q2tp_i8_dot_matches_exact_on_grid() {
10013        // On-grid activations (±1 → sx=1/127, xq=±127 dequantizes
10014        // exactly, no outliers) must make the integer path agree with
10015        // the exact scalar walk to f32 rounding.
10016        let (rows, cols) = (5, 64);
10017        let gpr = cols / GROUP_SIZE;
10018        // Synthetic codes plane + a flat ladder: scales_into is not under
10019        // test here, so drive dot_q2tp_row_i8 / q2tp_row_exact directly
10020        // with hand-made scales.
10021        let chunks: Vec<u8> = (0..rows * gpr * Q2TP_CHUNK)
10022            .map(|i| (i as u32).wrapping_mul(2654435761) as u8)
10023            .collect();
10024        let scales: Vec<f32> = (0..gpr).map(|g| 0.5 + g as f32 * 0.25).collect();
10025        let x: Vec<f32> = (0..cols)
10026            .map(|i| if i % 3 == 0 { -1.0 } else { 1.0 })
10027            .collect();
10028        let act = split_act(&x);
10029        assert!(
10030            act.outliers.is_empty(),
10031            "on-grid input must have no outliers"
10032        );
10033        let gsum = q1_group_sums(&act.xq, gpr);
10034        for r in 0..rows {
10035            let exact = q2tp_row_exact(&chunks, r, gpr, &x, &scales);
10036            let fast = dot_q2tp_row_i8(&chunks, r, gpr, &act.xq, &gsum, &scales) * act.sx;
10037            assert!(
10038                (exact - fast).abs() <= exact.abs() * 1e-5 + 1e-5,
10039                "row {r}: exact {exact} vs i8 {fast}"
10040            );
10041        }
10042    }
10043
10044    #[test]
10045    fn q8_row_dot_fast_matches_scalar() {
10046        // The per-arch fast dot must agree with the exact scalar oracle
10047        // (same contract the fused q8 FFN arm rides on).
10048        let cols = 96;
10049        let row: Vec<u8> = (0..cols)
10050            .map(|i| ((i * 37 % 251) - 125) as i8 as u8)
10051            .collect();
10052        let x: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.13).sin()).collect();
10053        let act = split_act(&x);
10054        let fast = q8_row_dot(&row, &act);
10055        let scalar = q8_row_dot_scalar(&row, &act);
10056        assert!(
10057            (fast - scalar).abs() <= scalar.abs() * 1e-5 + 1e-5,
10058            "fast {fast} vs scalar {scalar}"
10059        );
10060    }
10061
10062    #[test]
10063    fn f32_matvec_matches_matvec_rows_bitexact() {
10064        let (rows, cols) = (300, 40);
10065        let w: Vec<f32> = (0..rows * cols).map(|i| (i as f32 * 0.017).sin()).collect();
10066        let x: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.05).cos()).collect();
10067        let qt = QTensor::from_f32(w.clone(), rows, cols);
10068
10069        let mut a = vec![0.0f32; rows];
10070        matvec_rows(None, &w, &x, &mut a);
10071        let mut b = vec![0.0f32; rows];
10072        qt.matvec(&x, &mut b, None);
10073        assert_eq!(a, b);
10074    }
10075
10076    #[test]
10077    fn sdot_kernel_exact_on_grid() {
10078        // Activations already on the i8 grid (±1 with amax=1 → sx=1/127,
10079        // xq=±127 dequantizes EXACTLY) → the SDOT path must match the
10080        // exact f32 dot to float rounding. This isolates kernel
10081        // correctness from quantization noise.
10082        eprintln!("sdot_enabled = {}", sdot_enabled());
10083        let (rows, cols) = (9, 80); // odd rows → exercises 4-row + tail
10084        let w: Vec<u8> = (0..rows * cols)
10085            .map(|i| (((i * 37) % 251) as i32 - 125) as i8 as u8)
10086            .collect();
10087        let scales: Vec<f32> = (0..rows).map(|o| 0.005 + o as f32 * 0.001).collect();
10088        let x: Vec<f32> = (0..cols)
10089            .map(|i| match i % 3 {
10090                0 => 1.0,
10091                1 => -1.0,
10092                _ => 0.0,
10093            })
10094            .collect();
10095        let mut a = vec![0.0f32; rows];
10096        qmatvec(
10097            &w,
10098            &[],
10099            &scales,
10100            &x,
10101            &[],
10102            TensorDtype::Q8Row,
10103            rows,
10104            cols,
10105            &mut a,
10106            None,
10107        );
10108        for o in 0..rows {
10109            let mut acc = 0.0f32;
10110            for j in 0..cols {
10111                acc += (w[o * cols + j] as i8) as f32 * x[j];
10112            }
10113            let expect = acc * scales[o];
10114            assert!(
10115                (a[o] - expect).abs() < 1e-3 * expect.abs().max(1e-3),
10116                "row {o}: {} vs {expect}",
10117                a[o]
10118            );
10119        }
10120    }
10121
10122    #[test]
10123    fn q1_tbl_fast_path_matches_reference() {
10124        // gpr = 8 exercises the TBL pair-load fast loop, and the LAST
10125        // row's final 4-tile window trips the 4B-overread guard (the
10126        // payload ends exactly at the last tile) — both paths must
10127        // agree with the dequant reference.
10128        let (rows, cols) = (5, 256);
10129        let gpr = cols / GROUP_SIZE;
10130        let mut bytes = Vec::new();
10131        for t in 0..rows * gpr {
10132            let s = 0.007 + (t % 11) as f32 * 0.004;
10133            bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
10134            for j in 0..4 {
10135                bytes.push(((t * 53 + j * 89 + 7) % 249) as u8);
10136            }
10137        }
10138        let x: Vec<f32> = (0..cols)
10139            .map(|i| if (i * 5) % 7 < 3 { 1.0 } else { -1.0 })
10140            .collect();
10141        let mut w = vec![0.0f32; rows * cols];
10142        cortiq_core::quant::dequant_q1(&bytes, &mut w);
10143        let mut got = vec![0.0f32; rows];
10144        q1_matvec(&bytes, &x, rows, cols, &mut got, None);
10145        for o in 0..rows {
10146            let expect: f32 = (0..cols).map(|j| w[o * cols + j] * x[j]).sum();
10147            assert!(
10148                (got[o] - expect).abs() < 1e-3 * expect.abs().max(1e-3),
10149                "row {o}: {} vs {expect}",
10150                got[o]
10151            );
10152        }
10153        // Blocked 1×4 batch (b=5: one quad + remainder) must equal the
10154        // single-matvec path bit-for-bit.
10155        let b = 5usize;
10156        let mut xs_all = Vec::new();
10157        for bi in 0..b {
10158            xs_all.extend(x.iter().map(|v| if bi % 2 == 0 { *v } else { -*v }));
10159        }
10160        let mut mm = vec![0.0f32; b * rows];
10161        q1_matmat(&bytes, &xs_all, b, rows, cols, &mut mm, None);
10162        for bi in 0..b {
10163            let mut single = vec![0.0f32; rows];
10164            q1_matvec(
10165                &bytes,
10166                &xs_all[bi * cols..(bi + 1) * cols],
10167                rows,
10168                cols,
10169                &mut single,
10170                None,
10171            );
10172            assert_eq!(&mm[bi * rows..(bi + 1) * rows], &single[..], "stream {bi}");
10173        }
10174    }
10175
10176    #[test]
10177    fn q1_kernels_match_exact_reference() {
10178        // Synthetic q1 payload: 6-byte tiles [f16 scale][4B bits].
10179        let (rows, cols) = (7, 96);
10180        let gpr = cols / GROUP_SIZE;
10181        let mut bytes = Vec::new();
10182        for t in 0..rows * gpr {
10183            let s = 0.01 + (t % 13) as f32 * 0.003;
10184            bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
10185            for j in 0..4 {
10186                bytes.push(((t * 31 + j * 97) % 251) as u8);
10187            }
10188        }
10189        // On-grid activations (±1, amax 1) → the SDOT path is exact.
10190        let x: Vec<f32> = (0..cols)
10191            .map(|i| if i % 3 == 0 { 1.0 } else { -1.0 })
10192            .collect();
10193        // Reference through the core dequant.
10194        let mut w = vec![0.0f32; rows * cols];
10195        cortiq_core::quant::dequant_q1(&bytes, &mut w);
10196        let mut expect = vec![0.0f32; rows];
10197        for o in 0..rows {
10198            expect[o] = (0..cols).map(|j| w[o * cols + j] * x[j]).sum();
10199        }
10200        let mut got = vec![0.0f32; rows];
10201        q1_matvec(&bytes, &x, rows, cols, &mut got, None);
10202        for o in 0..rows {
10203            assert!(
10204                (got[o] - expect[o]).abs() < 1e-3 * expect[o].abs().max(1e-3),
10205                "row {o}: {} vs {}",
10206                got[o],
10207                expect[o]
10208            );
10209        }
10210        // Pair and batch paths agree with the single path.
10211        let x2: Vec<f32> = x.iter().map(|v| -v).collect();
10212        let (mut a1, mut a2) = (vec![0.0f32; rows], vec![0.0f32; rows]);
10213        q1_matvec2(&bytes, &x, &x2, rows, cols, &mut a1, &mut a2, None);
10214        assert_eq!(a1, got);
10215        let mut xs = x.clone();
10216        xs.extend_from_slice(&x2);
10217        let mut mm = vec![0.0f32; 2 * rows];
10218        q1_matmat(&bytes, &xs, 2, rows, cols, &mut mm, None);
10219        assert_eq!(&mm[..rows], got.as_slice());
10220        assert_eq!(&mm[rows..], a2.as_slice());
10221    }
10222
10223    #[test]
10224    fn repack_is_bit_identical() {
10225        // The interleaved-repack kernel must produce EXACTLY the same
10226        // bits as the mmap-layout kernel: integer accumulation is order-
10227        // exact, the f32 epilogue is identical. Odd rows exercise the
10228        // tail; direct range calls exercise unaligned pool splits.
10229        let (rows, cols) = (267, 96); // 66 groups + 3 tail rows, cols % 16 == 0
10230        let w: Vec<u8> = (0..rows * cols)
10231            .map(|i| (((i * 89) % 253) as i32 - 126) as i8 as u8)
10232            .collect();
10233        let scales: Vec<f32> = (0..rows).map(|o| 0.003 + o as f32 * 0.0007).collect();
10234        let x: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.37).sin() * 2.0).collect();
10235        let rep = q8_repack_layout(&w, rows, cols);
10236        // Group interleave round-trips.
10237        for g in 0..rows / 4 {
10238            for c in 0..cols / 16 {
10239                for lane in 0..4 {
10240                    assert_eq!(
10241                        &rep[g * 4 * cols + c * 64 + lane * 16
10242                            ..g * 4 * cols + c * 64 + lane * 16 + 16],
10243                        &w[(g * 4 + lane) * cols + c * 16..(g * 4 + lane) * cols + c * 16 + 16],
10244                    );
10245                }
10246            }
10247        }
10248        let mut a = vec![0.0f32; rows];
10249        qmatvec(
10250            &w,
10251            &[],
10252            &scales,
10253            &x,
10254            &[],
10255            TensorDtype::Q8Row,
10256            rows,
10257            cols,
10258            &mut a,
10259            None,
10260        );
10261        let mut b = vec![0.0f32; rows];
10262        qmatvec(
10263            &w,
10264            &rep,
10265            &scales,
10266            &x,
10267            &[],
10268            TensorDtype::Q8Row,
10269            rows,
10270            cols,
10271            &mut b,
10272            None,
10273        );
10274        assert_eq!(a, b, "full-range repack output diverged");
10275
10276        #[cfg(target_arch = "aarch64")]
10277        if sdot_enabled() {
10278            // Unaligned range split (pool workers get arbitrary bounds).
10279            let act = split_act(&x);
10280            let mut c1 = vec![0.0f32; rows];
10281            let mut c2 = vec![0.0f32; rows];
10282            q8_range_sdot(
10283                &w,
10284                &[],
10285                &scales,
10286                &act,
10287                cols,
10288                SendMut(c1.as_mut_ptr()),
10289                3,
10290                rows - 2,
10291            );
10292            q8_range_sdot(
10293                &w,
10294                &rep,
10295                &scales,
10296                &act,
10297                cols,
10298                SendMut(c2.as_mut_ptr()),
10299                3,
10300                rows - 2,
10301            );
10302            assert_eq!(c1, c2, "unaligned-range repack output diverged");
10303        }
10304    }
10305
10306    #[test]
10307    fn sdot_a8w8_noise_is_bounded() {
10308        // Off-grid activations: A8 quantization noise must stay small in
10309        // relative L2 over the whole output (realistic accuracy contract;
10310        // vmfcore measured argmax-identical decode on real models).
10311        let (rows, cols) = (16, 512);
10312        let w: Vec<u8> = (0..rows * cols)
10313            .map(|i| (((i * 37) % 251) as i32 - 125) as i8 as u8)
10314            .collect();
10315        let scales = vec![0.01f32; rows];
10316        let x: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.21).sin()).collect();
10317        let mut a = vec![0.0f32; rows];
10318        qmatvec(
10319            &w,
10320            &[],
10321            &scales,
10322            &x,
10323            &[],
10324            TensorDtype::Q8Row,
10325            rows,
10326            cols,
10327            &mut a,
10328            None,
10329        );
10330        let (mut num, mut den) = (0f64, 0f64);
10331        for o in 0..rows {
10332            let mut acc = 0.0f32;
10333            for j in 0..cols {
10334                acc += (w[o * cols + j] as i8) as f32 * x[j];
10335            }
10336            let expect = acc * scales[o];
10337            num += ((a[o] - expect) as f64).powi(2);
10338            den += (expect as f64).powi(2);
10339        }
10340        let rel = (num / den.max(1e-12)).sqrt();
10341        assert!(rel < 0.05, "A8W8 relative L2 error too high: {rel}");
10342    }
10343
10344    #[test]
10345    fn i8_dot_neon_matches_scalar() {
10346        let n = 100;
10347        let w: Vec<u8> = (0..n).map(|i| ((i * 37 + 11) % 251) as u8).collect();
10348        let x: Vec<f32> = (0..n).map(|i| (i as f32 * 0.13).sin()).collect();
10349        let mut scalar = 0.0f32;
10350        for j in 0..n {
10351            scalar += (w[j] as i8) as f32 * x[j];
10352        }
10353        let fast = dot_i8_f32(&w, &x);
10354        assert!((scalar - fast).abs() < 1e-3 * scalar.abs().max(1.0));
10355    }
10356
10357    /// Fused vbit matvec must match full dequant_vbit + dense matvec.
10358    #[test]
10359    fn vbitmatvec_matches_full_dequant() {
10360        let (rows, cols) = (6, 64);
10361        let ng = cols / GROUP_SIZE;
10362        // Hand-craft: bits per row, f16 scales, packed rows.
10363        let bits: Vec<u8> = vec![3, 4, 5, 6, 8, 4];
10364        let mut bytes = bits.clone();
10365        for g in 0..rows * ng {
10366            let s = 0.02 + 0.001 * g as f32;
10367            bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
10368        }
10369        for r in 0..rows {
10370            let b = bits[r] as usize;
10371            let (mut acc, mut nb) = (0u64, 0usize);
10372            let mut rowbytes = Vec::new();
10373            for i in 0..cols {
10374                let v = ((i * 7 + r * 13) % (1 << b)) as u64;
10375                acc = (acc << b) | v;
10376                nb += b;
10377                while nb >= 8 {
10378                    nb -= 8;
10379                    rowbytes.push(((acc >> nb) & 0xFF) as u8);
10380                }
10381            }
10382            if nb > 0 {
10383                rowbytes.push(((acc << (8 - nb)) & 0xFF) as u8);
10384            }
10385            bytes.extend_from_slice(&rowbytes);
10386        }
10387        let x: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.19).sin()).collect();
10388
10389        let mut reference = vec![0f32; rows * cols];
10390        cortiq_core::quant::dequant_vbit(&bytes, rows, cols, &mut reference).unwrap();
10391        let mut expect = vec![0f32; rows];
10392        for r in 0..rows {
10393            expect[r] = reference[r * cols..(r + 1) * cols]
10394                .iter()
10395                .zip(&x)
10396                .map(|(w, xv)| w * xv)
10397                .sum();
10398        }
10399        let mut got = vec![0f32; rows];
10400        let offsets = vbit_row_offsets(&bytes, rows, cols);
10401        vbitmatvec(&bytes, &offsets, &x, rows, cols, &mut got, None);
10402        // SDOT path quantizes activations to i8 (A8W8): bounded noise,
10403        // same contract as q8 (exact path is pinned by CMF_SDOT=0 in
10404        // the golden-parity gate).
10405        let tol = if a8w8_enabled() { 6e-2 } else { 1e-4 };
10406        let scale = expect.iter().fold(0f32, |m, v| m.max(v.abs())).max(1e-6);
10407        for r in 0..rows {
10408            assert!(
10409                (got[r] - expect[r]).abs() < tol * scale,
10410                "row {r}: {} vs {}",
10411                got[r],
10412                expect[r]
10413            );
10414        }
10415    }
10416
10417    /// Fused q4 matvec must match the reference full-dequant + dense
10418    /// matvec bit-for-bit in structure (same f32 math, group order).
10419    /// vbit matmat: the blocked 1×4 leg must match the per-row path
10420    /// (paired env toggle; larger shape so both code paths engage).
10421    #[test]
10422    #[cfg(target_arch = "x86_64")]
10423    fn vbit_matmat_blocked_matches_per_row() {
10424        let (rows, cols, b) = (64usize, 128usize, 9usize);
10425        let ng = cols / GROUP_SIZE;
10426        let bits: Vec<u8> = (0..rows).map(|r| [3u8, 4, 5, 6][r % 4]).collect();
10427        let mut bytes = bits.clone();
10428        for g in 0..rows * ng {
10429            let sc = 0.02 + 0.0005 * g as f32;
10430            bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(sc).to_le_bytes());
10431        }
10432        for r in 0..rows {
10433            let bw = bits[r] as usize;
10434            let (mut acc, mut nb) = (0u64, 0usize);
10435            let mut rowbytes = Vec::new();
10436            for i in 0..cols {
10437                let v = ((i * 7 + r * 13) % (1 << bw)) as u64;
10438                acc = (acc << bw) | v;
10439                nb += bw;
10440                while nb >= 8 {
10441                    nb -= 8;
10442                    rowbytes.push(((acc >> nb) & 0xFF) as u8);
10443                }
10444            }
10445            if nb > 0 {
10446                rowbytes.push(((acc << (8 - nb)) & 0xFF) as u8);
10447            }
10448            bytes.extend_from_slice(&rowbytes);
10449        }
10450        let x: Vec<f32> = (0..b * cols)
10451            .map(|i| ((i * 13 + 7) % 97) as f32 / 97.0 - 0.5)
10452            .collect();
10453        let offsets = vbit_row_offsets(&bytes, rows, cols);
10454        let mut y_a = vec![0f32; b * rows];
10455        let mut y_b = vec![0f32; b * rows];
10456        unsafe { std::env::set_var("CMF_X86_BLOCKED", "1") };
10457        vbitmatmat(&bytes, &offsets, &x, b, rows, cols, &mut y_a, None);
10458        unsafe { std::env::set_var("CMF_X86_BLOCKED", "0") };
10459        vbitmatmat(&bytes, &offsets, &x, b, rows, cols, &mut y_b, None);
10460        unsafe { std::env::remove_var("CMF_X86_BLOCKED") };
10461        let max_d = y_a
10462            .iter()
10463            .zip(&y_b)
10464            .map(|(p, q)| (p - q).abs())
10465            .fold(0.0f32, f32::max);
10466        assert!(max_d < 1e-4, "vbit blocked ≠ per-row: max|Δ| = {max_d}");
10467    }
10468
10469    /// q4t blocked 1×4 (SDOT on ARM, AVX2 on x86) must equal the
10470    /// per-row path exactly: same nibble unpack, same group order,
10471    /// same f32 accumulation — batch == matvec bit-for-bit. b=9 covers
10472    /// two full 1×4 blocks plus a remainder through the single-row
10473    /// kernel. (Both paths produce identical output, so the shared
10474    /// CMF_X86_BLOCKED env var racing with other tests cannot flip
10475    /// the verdict — worst case both sides take the same path.)
10476    #[test]
10477    fn q4t_matmat_blocked_matches_per_row() {
10478        let (rows, cols, b) = (16usize, 64usize, 9usize);
10479        let gpr = cols / GROUP_SIZE;
10480        let mut bytes = vec![0u8; rows * gpr * Q4_TILE];
10481        for r in 0..rows {
10482            for g in 0..gpr {
10483                let t = (r * gpr + g) * Q4_TILE;
10484                let sc = 0.02 + 0.001 * (r * gpr + g) as f32;
10485                bytes[t..t + 2].copy_from_slice(&cortiq_core::quant::f32_to_f16(sc).to_le_bytes());
10486                for k in 0..16 {
10487                    bytes[t + 2 + k] = ((r * 31 + g * 7 + k * 13) % 251) as u8;
10488                }
10489            }
10490        }
10491        let x: Vec<f32> = (0..b * cols)
10492            .map(|i| ((i * 13 + 7) % 97) as f32 / 97.0 - 0.5)
10493            .collect();
10494        let mut y_blk = vec![0f32; b * rows];
10495        let mut y_row = vec![0f32; b * rows];
10496        unsafe { std::env::set_var("CMF_X86_BLOCKED", "1") };
10497        q4t_matmat(&bytes, &x, b, rows, cols, &mut y_blk, None);
10498        unsafe { std::env::set_var("CMF_X86_BLOCKED", "0") };
10499        q4t_matmat(&bytes, &x, b, rows, cols, &mut y_row, None);
10500        unsafe { std::env::remove_var("CMF_X86_BLOCKED") };
10501        assert_eq!(y_blk, y_row, "q4t blocked 1x4 ≠ per-row");
10502    }
10503
10504    /// The wide-batch Accelerate arm of q4t_matmat vs a brute-force
10505    /// f32 dequant matmul: both are f32 GEMMs, so only reduction
10506    /// order differs — tight tolerance.
10507    /// A synthetic q4tp payload: random nibbles plus a per-row ladder whose
10508    /// span varies row to row, so the codes actually exercise the full 0..31
10509    /// range rather than clustering on one rung.
10510    fn synth_q4tp(rows: usize, cols: usize) -> Vec<u8> {
10511        use cortiq_core::quant::{f32_to_f16, q4tp_code_stride, q4tp_put_code};
10512        let gpr = cols / GROUP_SIZE;
10513        let stride = q4tp_code_stride(gpr);
10514        let (params_off, codes_off, _) = q4tp_sections(rows, cols);
10515        let mut b = vec![0u8; codes_off + rows * stride];
10516        for r in 0..rows {
10517            for g in 0..gpr {
10518                let t = (r * gpr + g) * Q4TP_NIB;
10519                for k in 0..16 {
10520                    b[t + k] = ((r * 31 + g * 7 + k * 13) % 251) as u8;
10521                }
10522            }
10523            let lo = -6.0 - 0.03 * (r % 17) as f32;
10524            let step = 0.01 + 0.004 * (r % 11) as f32;
10525            let p = params_off + r * 4;
10526            b[p..p + 2].copy_from_slice(&f32_to_f16(lo).to_le_bytes());
10527            b[p + 2..p + 4].copy_from_slice(&f32_to_f16(step).to_le_bytes());
10528            let crow = &mut b[codes_off + r * stride..codes_off + (r + 1) * stride];
10529            for g in 0..gpr {
10530                q4tp_put_code(crow, g, (r * 5 + g * 3) % 32);
10531            }
10532        }
10533        b
10534    }
10535
10536    /// The same weights re-expressed as q4_tiled, so the proven kernel can
10537    /// be the reference: each tile stores the ladder scale its code selects.
10538    /// Only the f16 rounding of that scale separates the two payloads.
10539    fn q4tp_as_q4t(bytes: &[u8], rows: usize, cols: usize) -> Vec<u8> {
10540        let gpr = cols / GROUP_SIZE;
10541        let v = Q4tpView::new(bytes, rows, cols);
10542        let mut out = vec![0u8; rows * gpr * Q4_TILE];
10543        let mut sc = vec![0f32; gpr];
10544        for r in 0..rows {
10545            v.scales_into(r, gpr, &mut sc);
10546            for g in 0..gpr {
10547                let t = (r * gpr + g) * Q4_TILE;
10548                let s = sc[g];
10549                out[t..t + 2].copy_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
10550                let src = (r * gpr + g) * Q4TP_NIB;
10551                out[t + 2..t + Q4_TILE].copy_from_slice(&v.nib[src..src + Q4TP_NIB]);
10552            }
10553        }
10554        out
10555    }
10556
10557    /// The exact (`CMF_SDOT=0`) path must reproduce `dequant_q4tp` to f32
10558    /// rounding — that scalar routine is the format's definition, and the
10559    /// kernels re-derive the scale from the ladder independently. Call the
10560    /// row kernel directly: `matmat` picks the int8 arm when a8w8 is on,
10561    /// so routing through it would test the other path by accident.
10562    #[test]
10563    fn q4tp_exact_path_matches_dequant_reference() {
10564        let (rows, cols) = (256usize, 512usize);
10565        let gpr = cols / GROUP_SIZE;
10566        let bytes = synth_q4tp(rows, cols);
10567        let mut w = vec![0f32; rows * cols];
10568        cortiq_core::quant::dequant_q4tp(&bytes, rows, cols, &mut w);
10569
10570        let x: Vec<f32> = (0..cols)
10571            .map(|i| ((i * 13 + 7) % 97) as f32 / 97.0 - 0.5)
10572            .collect();
10573        let v = Q4tpView::new(&bytes, rows, cols);
10574        let mut sc = vec![0f32; gpr];
10575        for r in 0..rows {
10576            v.scales_into(r, gpr, &mut sc);
10577            let got = q4tp_row_exact(v.nib, r, gpr, &x, &sc);
10578            let want: f32 = (0..cols).map(|c| w[r * cols + c] * x[c]).sum();
10579            // These dot products cancel down to ~1e-3 from terms of ~5e-2, so
10580            // the meaningful yardstick is the summed magnitude, not the result:
10581            // against the result any reordering of a 512-term f32 sum "fails".
10582            let mag: f32 = (0..cols).map(|c| (w[r * cols + c] * x[c]).abs()).sum();
10583            assert!(
10584                (got - want).abs() <= 1e-5 * mag,
10585                "row {r}: kernel {got} vs dequant {want}"
10586            );
10587        }
10588    }
10589
10590    /// The int8 (a8w8) path can't be checked against an f32 reference — the
10591    /// activation quantization dominates. Check it against the q4t kernel it
10592    /// was ported from instead, on payloads holding the same weights: that
10593    /// isolates exactly what the port could break (16 B stride, ladder
10594    /// lookup, nibble unpack) from what it deliberately shares.
10595    #[test]
10596    fn q4tp_matvec_matches_the_q4t_kernel_it_was_ported_from() {
10597        let (rows, cols) = (256usize, 512usize);
10598        let bytes = synth_q4tp(rows, cols);
10599        let twin = q4tp_as_q4t(&bytes, rows, cols);
10600        let x: Vec<f32> = (0..cols)
10601            .map(|i| ((i * 13 + 7) % 97) as f32 / 97.0 - 0.5)
10602            .collect();
10603
10604        let mut got = vec![0f32; rows];
10605        q4tp_matvec(&bytes, &x, rows, cols, &mut got, None);
10606        let mut want = vec![0f32; rows];
10607        q4t_matvec(&twin, &x, rows, cols, &mut want, None);
10608
10609        // Scale is f16 in the twin and f32 here, so allow that rounding on
10610        // top of the summed magnitude (same cancellation argument as above).
10611        let mut w = vec![0f32; rows * cols];
10612        cortiq_core::quant::dequant_q4tp(&bytes, rows, cols, &mut w);
10613        for r in 0..rows {
10614            let mag: f32 = (0..cols).map(|c| (w[r * cols + c] * x[c]).abs()).sum();
10615            assert!(
10616                (got[r] - want[r]).abs() <= 1e-3 * mag,
10617                "row {r}: q4tp {} vs q4t {}",
10618                got[r],
10619                want[r]
10620            );
10621        }
10622    }
10623
10624    /// `matmat` carries three arms (Accelerate, blocked int8 1x4, scalar).
10625    /// Batch 5 crosses the blocked kernel's stride, so this exercises the
10626    /// 1x4 path AND its scalar tail in one run — the blocked kernel is new
10627    /// code and its four accumulators are exactly what tends to go wrong.
10628    #[test]
10629    fn q4tp_matmat_matches_the_q4t_kernel_it_was_ported_from() {
10630        let (rows, cols, b) = (256usize, 512usize, 5usize);
10631        let bytes = synth_q4tp(rows, cols);
10632        let twin = q4tp_as_q4t(&bytes, rows, cols);
10633        let xs: Vec<f32> = (0..b * cols)
10634            .map(|i| ((i * 29 + 11) % 89) as f32 / 89.0 - 0.5)
10635            .collect();
10636
10637        let mut got = vec![0f32; b * rows];
10638        q4tp_matmat(&bytes, &xs, b, rows, cols, &mut got, None);
10639        let mut want = vec![0f32; b * rows];
10640        q4t_matmat(&twin, &xs, b, rows, cols, &mut want, None);
10641
10642        let mut w = vec![0f32; rows * cols];
10643        cortiq_core::quant::dequant_q4tp(&bytes, rows, cols, &mut w);
10644        for t in 0..b {
10645            for r in 0..rows {
10646                let mag: f32 = (0..cols)
10647                    .map(|c| (w[r * cols + c] * xs[t * cols + c]).abs())
10648                    .sum();
10649                let (g, wa) = (got[t * rows + r], want[t * rows + r]);
10650                assert!(
10651                    (g - wa).abs() <= 1e-3 * mag,
10652                    "batch {t} row {r}: q4tp {g} vs q4t {wa}"
10653                );
10654            }
10655        }
10656    }
10657
10658    #[test]
10659    fn q4tp_matvec2_matches_the_single_stream_kernel() {
10660        let (rows, cols) = (128usize, 256usize);
10661        let gpr = cols / GROUP_SIZE;
10662        let bytes = synth_q4tp(rows, cols);
10663        let xs: Vec<f32> = (0..2 * cols)
10664            .map(|i| ((i * 29 + 11) % 89) as f32 / 89.0 - 0.5)
10665            .collect();
10666
10667        let (mut o1, mut o2) = (vec![0f32; rows], vec![0f32; rows]);
10668        q4tp_matvec2(
10669            &bytes,
10670            &xs[..cols],
10671            &xs[cols..],
10672            rows,
10673            cols,
10674            &mut o1,
10675            &mut o2,
10676            None,
10677        );
10678
10679        // matvec2 takes the exact path for both streams, so the single-row
10680        // kernel is an exact reference — no tolerance for path differences.
10681        let v = Q4tpView::new(&bytes, rows, cols);
10682        let mut sc = vec![0f32; gpr];
10683        for r in 0..rows {
10684            v.scales_into(r, gpr, &mut sc);
10685            assert_eq!(o1[r], q4tp_row_exact(v.nib, r, gpr, &xs[..cols], &sc));
10686            assert_eq!(o2[r], q4tp_row_exact(v.nib, r, gpr, &xs[cols..], &sc));
10687        }
10688    }
10689
10690    /// q4tp must not COST speed — it exists to save bytes, and a format that
10691    /// trades 7% of a file for a slower model is a bad trade. This guard is
10692    /// here because correctness tests happily passed while `q4tp_matmat` was
10693    /// missing its int8 and Accelerate arms and the model ran 5x slower.
10694    /// Measured on M-series: 0.97-1.04x, i.e. parity (16 B tiles are better
10695    /// aligned than q4t's 18 B, which pays for the scale indirection).
10696    #[test]
10697    fn q4tp_matvec_keeps_pace_with_q4t() {
10698        let (rows, cols) = (4096usize, 3072usize);
10699        let bytes = synth_q4tp(rows, cols);
10700        let twin = q4tp_as_q4t(&bytes, rows, cols);
10701        let x: Vec<f32> = (0..cols).map(|i| (i % 97) as f32 / 97.0 - 0.5).collect();
10702        let mut o = vec![0f32; rows];
10703        let n = 12;
10704        let mut best = (f64::MAX, f64::MAX);
10705        // Interleaved A/B, minimum statistic: this machine throttles, and a
10706        // mean over a thermal ramp reliably indicts whichever ran second.
10707        for _ in 0..3 {
10708            let t0 = std::time::Instant::now();
10709            for _ in 0..n {
10710                q4t_matvec(&twin, &x, rows, cols, &mut o, None);
10711            }
10712            best.0 = best.0.min(t0.elapsed().as_secs_f64());
10713            let t0 = std::time::Instant::now();
10714            for _ in 0..n {
10715                q4tp_matvec(&bytes, &x, rows, cols, &mut o, None);
10716            }
10717            best.1 = best.1.min(t0.elapsed().as_secs_f64());
10718        }
10719        let ratio = best.1 / best.0;
10720        println!(
10721            "q4t {:.3} ms | q4tp {:.3} ms | {ratio:.2}x",
10722            best.0 * 1e3 / n as f64,
10723            best.1 * 1e3 / n as f64
10724        );
10725        assert!(ratio < 2.0, "q4tp matvec {ratio:.2}x slower than q4t");
10726    }
10727
10728    #[cfg(target_os = "macos")]
10729    #[test]
10730    fn q4t_matmat_accel_matches_dequant_reference() {
10731        if !accel_gemm_enabled() {
10732            return; // CMF_ACCEL=0
10733        }
10734        let (rows, cols, b) = (512usize, 1024usize, 8usize); // ≥500K → accel arm
10735        let gpr = cols / GROUP_SIZE;
10736        let mut bytes = vec![0u8; rows * gpr * Q4_TILE];
10737        for r in 0..rows {
10738            for g in 0..gpr {
10739                let t = (r * gpr + g) * Q4_TILE;
10740                let sc = 0.02 + 0.0005 * ((r * gpr + g) % 64) as f32;
10741                bytes[t..t + 2].copy_from_slice(&cortiq_core::quant::f32_to_f16(sc).to_le_bytes());
10742                for k in 0..16 {
10743                    bytes[t + 2 + k] = ((r * 31 + g * 7 + k * 13) % 251) as u8;
10744                }
10745            }
10746        }
10747        let x: Vec<f32> = (0..b * cols)
10748            .map(|i| ((i * 13 + 7) % 97) as f32 / 97.0 - 0.5)
10749            .collect();
10750        let mut got = vec![0f32; b * rows];
10751        q4t_matmat(&bytes, &x, b, rows, cols, &mut got, None);
10752        // Brute-force reference off the same tiles.
10753        let mut w = vec![0f32; rows * cols];
10754        for r in 0..rows {
10755            for g in 0..gpr {
10756                let t = (r * gpr + g) * Q4_TILE;
10757                let s = f16_to_f32(u16::from_le_bytes([bytes[t], bytes[t + 1]]));
10758                for (k, &bb) in bytes[t + 2..t + Q4_TILE].iter().enumerate() {
10759                    w[r * cols + g * GROUP_SIZE + k * 2] = ((bb & 0x0F) as f32 - 8.0) * s;
10760                    w[r * cols + g * GROUP_SIZE + k * 2 + 1] =
10761                        (((bb >> 4) & 0x0F) as f32 - 8.0) * s;
10762                }
10763            }
10764        }
10765        for bi in 0..b {
10766            for r in 0..rows {
10767                let want: f32 = (0..cols).map(|j| x[bi * cols + j] * w[r * cols + j]).sum();
10768                let d = (got[bi * rows + r] - want).abs();
10769                assert!(
10770                    d <= want.abs().max(1.0) * 1e-4,
10771                    "accel q4t GEMM diverged at ({bi},{r}): {} vs {want}",
10772                    got[bi * rows + r]
10773                );
10774            }
10775        }
10776    }
10777
10778    #[test]
10779    fn q4matvec_matches_full_dequant() {
10780        let (rows, cols) = (8, 64);
10781        let groups = rows * cols / GROUP_SIZE;
10782        // Hand-craft a q4_block blob: nibbles then f16 scales.
10783        let mut bytes = Vec::with_capacity(groups * 16 + groups * 2);
10784        for i in 0..groups * 16 {
10785            bytes.push((((i * 7 + 3) % 256) & 0xFF) as u8);
10786        }
10787        for g in 0..groups {
10788            let s = 0.01 + 0.003 * g as f32;
10789            bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
10790        }
10791        let x: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.17).sin()).collect();
10792
10793        let mut reference = vec![0.0f32; rows * cols];
10794        cortiq_core::quant::dequant_q4_block(&bytes, &mut reference);
10795        let mut expect = vec![0.0f32; rows];
10796        for r in 0..rows {
10797            expect[r] = reference[r * cols..(r + 1) * cols]
10798                .iter()
10799                .zip(&x)
10800                .map(|(w, xv)| w * xv)
10801                .sum();
10802        }
10803
10804        let mut got = vec![0.0f32; rows];
10805        q4matvec(&bytes, &x, rows, cols, &mut got, None);
10806        // SDOT path quantizes activations to i8 (A8W8): bounded noise,
10807        // same contract as q8/vbit (exact path is pinned by CMF_SDOT=0
10808        // in the golden-parity gate).
10809        let tol = if a8w8_enabled() { 6e-2 } else { 1e-4 };
10810        let scale = expect.iter().fold(0f32, |m, v| m.max(v.abs())).max(1.0);
10811        for r in 0..rows {
10812            assert!(
10813                (got[r] - expect[r]).abs() < tol * scale,
10814                "row {r}: {} vs {}",
10815                got[r],
10816                expect[r]
10817            );
10818        }
10819    }
10820
10821    /// Fused two-input vbit matvec must equal two single matvecs exactly
10822    /// (same per-lane accumulation order on both scalar and SDOT paths).
10823    #[test]
10824    fn vbitmatvec2_equals_two_singles() {
10825        let (rows, cols) = (6, 64);
10826        let ng = cols / GROUP_SIZE;
10827        let bits: Vec<u8> = vec![3, 4, 5, 6, 8, 4];
10828        let mut bytes = bits.clone();
10829        for g in 0..rows * ng {
10830            let s = 0.02 + 0.001 * g as f32;
10831            bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
10832        }
10833        for r in 0..rows {
10834            let b = bits[r] as usize;
10835            let (mut acc, mut nb) = (0u64, 0usize);
10836            let mut rowbytes = Vec::new();
10837            for i in 0..cols {
10838                let v = ((i * 7 + r * 13) % (1 << b)) as u64;
10839                acc = (acc << b) | v;
10840                nb += b;
10841                while nb >= 8 {
10842                    nb -= 8;
10843                    rowbytes.push(((acc >> nb) & 0xFF) as u8);
10844                }
10845            }
10846            if nb > 0 {
10847                rowbytes.push(((acc << (8 - nb)) & 0xFF) as u8);
10848            }
10849            bytes.extend_from_slice(&rowbytes);
10850        }
10851        let x1: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.19).sin()).collect();
10852        let x2: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.11).cos()).collect();
10853        let offsets = vbit_row_offsets(&bytes, rows, cols);
10854
10855        let (mut a1, mut a2) = (vec![0f32; rows], vec![0f32; rows]);
10856        vbitmatvec(&bytes, &offsets, &x1, rows, cols, &mut a1, None);
10857        vbitmatvec(&bytes, &offsets, &x2, rows, cols, &mut a2, None);
10858        let (mut b1, mut b2) = (vec![0f32; rows], vec![0f32; rows]);
10859        vbitmatvec2(
10860            &bytes, &offsets, &x1, &x2, rows, cols, &mut b1, &mut b2, None,
10861        );
10862        assert_eq!(a1, b1, "fused vbit lane 1 must be bit-identical");
10863        assert_eq!(a2, b2, "fused vbit lane 2 must be bit-identical");
10864    }
10865
10866    /// Fused two-input q4 matvec must equal two single matvecs exactly.
10867    #[test]
10868    fn q4matvec2_equals_two_singles() {
10869        let (rows, cols) = (8, 128);
10870        let groups = rows * cols / GROUP_SIZE;
10871        let mut bytes = Vec::with_capacity(groups * 16 + groups * 2);
10872        for i in 0..groups * 16 {
10873            bytes.push((((i * 7 + 3) % 256) & 0xFF) as u8);
10874        }
10875        for g in 0..groups {
10876            let s = 0.01 + 0.003 * g as f32;
10877            bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
10878        }
10879        // Include an outlier channel so the SDOT correction path is
10880        // exercised in the pair kernel too.
10881        let mut x1: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.17).sin()).collect();
10882        x1[9] = 250.0;
10883        let x2: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.23).cos()).collect();
10884
10885        let (mut a1, mut a2) = (vec![0f32; rows], vec![0f32; rows]);
10886        q4matvec(&bytes, &x1, rows, cols, &mut a1, None);
10887        q4matvec(&bytes, &x2, rows, cols, &mut a2, None);
10888        let (mut b1, mut b2) = (vec![0f32; rows], vec![0f32; rows]);
10889        q4matvec2(&bytes, &x1, &x2, rows, cols, &mut b1, &mut b2, None);
10890        assert_eq!(a1, b1, "fused q4 lane 1 must be bit-identical");
10891        assert_eq!(a2, b2, "fused q4 lane 2 must be bit-identical");
10892    }
10893
10894    /// Multi-matrix job must equal separate matvecs exactly — same
10895    /// kernels, only the dispatch is fused.
10896    #[test]
10897    fn matvec_many_equals_separate_matvecs() {
10898        use crate::pool::Pool;
10899        let (r1, r2, cols) = (300, 200, 64);
10900        let mk = |salt: usize, rows: usize| {
10901            QTensor::from_f32(
10902                (0..rows * cols)
10903                    .map(|i| ((i * 7 + salt) % 97) as f32 / 97.0 - 0.5)
10904                    .collect(),
10905                rows,
10906                cols,
10907            )
10908        };
10909        let (a, b) = (mk(1, r1), mk(5, r2));
10910        let x: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.11).sin()).collect();
10911        let pool = Pool::new(3);
10912
10913        let (mut ea, mut eb) = (vec![0f32; r1], vec![0f32; r2]);
10914        a.matvec(&x, &mut ea, Some(&pool));
10915        b.matvec(&x, &mut eb, Some(&pool));
10916        let (mut ga, mut gb) = (vec![0f32; r1], vec![0f32; r2]);
10917        QTensor::matvec_many([&a, &b], &x, [&mut ga, &mut gb], Some(&pool));
10918        assert_eq!(ea, ga, "fused multi-matrix lane 1 must be bit-identical");
10919        assert_eq!(eb, gb, "fused multi-matrix lane 2 must be bit-identical");
10920    }
10921
10922    /// Batched q4/vbit matmat must equal per-position matvec calls
10923    /// exactly (the fallback it replaced) — same kernels, same order.
10924    #[test]
10925    fn batched_matmat_equals_per_position_matvec() {
10926        let (rows, cols, b) = (8, 64, 5);
10927        // q4 blob.
10928        let groups = rows * cols / GROUP_SIZE;
10929        let mut q4 = Vec::new();
10930        for i in 0..groups * 16 {
10931            q4.push((((i * 7 + 3) % 256) & 0xFF) as u8);
10932        }
10933        for g in 0..groups {
10934            q4.extend_from_slice(
10935                &cortiq_core::quant::f32_to_f16(0.01 + 0.003 * g as f32).to_le_bytes(),
10936            );
10937        }
10938        // vbit blob (mixed widths incl. 8).
10939        let ng = cols / GROUP_SIZE;
10940        let bits: Vec<u8> = vec![3, 4, 5, 6, 8, 4, 5, 3];
10941        let mut vb = bits.clone();
10942        for g in 0..rows * ng {
10943            vb.extend_from_slice(
10944                &cortiq_core::quant::f32_to_f16(0.02 + 0.001 * g as f32).to_le_bytes(),
10945            );
10946        }
10947        for r in 0..rows {
10948            let bw = bits[r] as usize;
10949            let (mut acc, mut nb) = (0u64, 0usize);
10950            let mut rowbytes = Vec::new();
10951            for i in 0..cols {
10952                let v = ((i * 7 + r * 13) % (1 << bw)) as u64;
10953                acc = (acc << bw) | v;
10954                nb += bw;
10955                while nb >= 8 {
10956                    nb -= 8;
10957                    rowbytes.push(((acc >> nb) & 0xFF) as u8);
10958                }
10959            }
10960            if nb > 0 {
10961                rowbytes.push(((acc << (8 - nb)) & 0xFF) as u8);
10962            }
10963            vb.extend_from_slice(&rowbytes);
10964        }
10965        let offsets = vbit_row_offsets(&vb, rows, cols);
10966
10967        let xs: Vec<f32> = (0..b * cols).map(|i| (i as f32 * 0.13).sin()).collect();
10968
10969        // q4: batch vs singles.
10970        let mut got = vec![0f32; b * rows];
10971        q4matmat(&q4, &xs, b, rows, cols, &mut got, None);
10972        for bi in 0..b {
10973            let mut expect = vec![0f32; rows];
10974            q4matvec(
10975                &q4,
10976                &xs[bi * cols..(bi + 1) * cols],
10977                rows,
10978                cols,
10979                &mut expect,
10980                None,
10981            );
10982            assert_eq!(
10983                &got[bi * rows..(bi + 1) * rows],
10984                &expect[..],
10985                "q4 batch pos {bi}"
10986            );
10987        }
10988
10989        // vbit: batch vs singles.
10990        let mut got = vec![0f32; b * rows];
10991        vbitmatmat(&vb, &offsets, &xs, b, rows, cols, &mut got, None);
10992        for bi in 0..b {
10993            let mut expect = vec![0f32; rows];
10994            vbitmatvec(
10995                &vb,
10996                &offsets,
10997                &xs[bi * cols..(bi + 1) * cols],
10998                rows,
10999                cols,
11000                &mut expect,
11001                None,
11002            );
11003            assert_eq!(
11004                &got[bi * rows..(bi + 1) * rows],
11005                &expect[..],
11006                "vbit batch pos {bi}"
11007            );
11008        }
11009    }
11010
11011    /// q4_tiled kernels must produce BIT-identical outputs to the q4
11012    /// split kernels on the same values (same ints, same order — only
11013    /// the byte placement differs).
11014    #[test]
11015    fn q4_tiled_matches_q4_block_bitexact() {
11016        let (rows, cols, b) = (8usize, 128usize, 3usize);
11017        let groups = rows * cols / GROUP_SIZE;
11018        let mut split = Vec::with_capacity(groups * 18);
11019        for i in 0..groups * 16 {
11020            split.push((((i * 7 + 3) % 256) & 0xFF) as u8);
11021        }
11022        for g in 0..groups {
11023            split.extend_from_slice(
11024                &cortiq_core::quant::f32_to_f16(0.01 + 0.003 * g as f32).to_le_bytes(),
11025            );
11026        }
11027        // Re-tile: [scale][nibbles] per group.
11028        let (packed, scales) = split.split_at(groups * 16);
11029        let mut tiled = Vec::with_capacity(groups * Q4_TILE);
11030        for g in 0..groups {
11031            tiled.extend_from_slice(&scales[g * 2..g * 2 + 2]);
11032            tiled.extend_from_slice(&packed[g * 16..(g + 1) * 16]);
11033        }
11034
11035        let mut x1: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.17).sin()).collect();
11036        x1[9] = 250.0; // exercise the outlier path
11037        let x2: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.23).cos()).collect();
11038
11039        let (mut a, mut t) = (vec![0f32; rows], vec![0f32; rows]);
11040        q4matvec(&split, &x1, rows, cols, &mut a, None);
11041        q4t_matvec(&tiled, &x1, rows, cols, &mut t, None);
11042        assert_eq!(a, t, "q4t matvec must match q4 bit-for-bit");
11043
11044        let (mut a1, mut a2) = (vec![0f32; rows], vec![0f32; rows]);
11045        let (mut t1, mut t2) = (vec![0f32; rows], vec![0f32; rows]);
11046        q4matvec2(&split, &x1, &x2, rows, cols, &mut a1, &mut a2, None);
11047        q4t_matvec2(&tiled, &x1, &x2, rows, cols, &mut t1, &mut t2, None);
11048        assert_eq!(a1, t1);
11049        assert_eq!(a2, t2);
11050
11051        let xs: Vec<f32> = (0..b * cols).map(|i| (i as f32 * 0.13).sin()).collect();
11052        let (mut am, mut tm) = (vec![0f32; b * rows], vec![0f32; b * rows]);
11053        q4matmat(&split, &xs, b, rows, cols, &mut am, None);
11054        q4t_matmat(&tiled, &xs, b, rows, cols, &mut tm, None);
11055        assert_eq!(am, tm, "q4t matmat must match q4 bit-for-bit");
11056    }
11057
11058    /// q4 SDOT outlier correction: a single huge activation channel
11059    /// (>8·rms → outlier, zeroed in xq) must still contribute its EXACT
11060    /// term. On-grid bulk (±1/0 → xq dequantizes exactly) isolates the
11061    /// correction from A8W8 noise. cols must exceed 64: at n=64 the
11062    /// 8·rms threshold equals sqrt(v²+rest) ≥ v, so a single outlier
11063    /// can never qualify (8² = n).
11064    #[test]
11065    fn q4matvec_sdot_outlier_exact() {
11066        let (rows, cols) = (4, 128);
11067        let groups = rows * cols / GROUP_SIZE;
11068        let mut bytes = Vec::with_capacity(groups * 16 + groups * 2);
11069        for i in 0..groups * 16 {
11070            bytes.push(((i * 11 + 5) % 256) as u8);
11071        }
11072        for g in 0..groups {
11073            let s = 0.02 + 0.002 * g as f32;
11074            bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
11075        }
11076        let mut x: Vec<f32> = (0..cols)
11077            .map(|i| match i % 3 {
11078                0 => 1.0,
11079                1 => -1.0,
11080                _ => 0.0,
11081            })
11082            .collect();
11083        x[17] = 300.0; // ≫ 8·rms → outlier channel
11084
11085        let mut reference = vec![0.0f32; rows * cols];
11086        cortiq_core::quant::dequant_q4_block(&bytes, &mut reference);
11087        let mut expect = vec![0.0f32; rows];
11088        for r in 0..rows {
11089            expect[r] = reference[r * cols..(r + 1) * cols]
11090                .iter()
11091                .zip(&x)
11092                .map(|(w, xv)| w * xv)
11093                .sum();
11094        }
11095        let mut got = vec![0.0f32; rows];
11096        q4matvec(&bytes, &x, rows, cols, &mut got, None);
11097        let scale = expect.iter().fold(0f32, |m, v| m.max(v.abs())).max(1.0);
11098        for r in 0..rows {
11099            assert!(
11100                (got[r] - expect[r]).abs() < 2e-3 * scale,
11101                "row {r}: {} vs {} (outlier term must be exact)",
11102                got[r],
11103                expect[r]
11104            );
11105        }
11106    }
11107
11108    /// The fused q1t matvec must equal the reference (dequant_q1t → dot),
11109    /// including the ternary zero level and the binary-searched outlier
11110    /// overlay. Guards the mmap kernel that makes a 12B q1t runnable.
11111    #[test]
11112    fn q1t_matvec_matches_reference() {
11113        use cortiq_core::quant::{dequant_q1t, f32_to_f16};
11114        let (rows, cols) = (3usize, 64usize); // gpr = 2
11115        let gpr = cols / GROUP_SIZE;
11116        let scales = [0.5f32, 0.3, 0.7, 0.2, 0.6, 0.15];
11117        // Overlay (must be sorted by flat index): a few spikes across rows.
11118        let outliers: [(u32, f32); 3] = [(5, 9.0), (70, -4.5), (150, 3.25)];
11119        let is_out = |flat: usize| outliers.iter().any(|&(i, _)| i as usize == flat);
11120        let mut bytes = Vec::new();
11121        for r in 0..rows {
11122            for g in 0..gpr {
11123                bytes.extend_from_slice(&f32_to_f16(scales[r * gpr + g]).to_le_bytes());
11124                let mut c = [0u8; 7];
11125                for k in 0..GROUP_SIZE {
11126                    // Encoder invariant: code 0 at outlier positions.
11127                    let code = if is_out(r * cols + g * GROUP_SIZE + k) {
11128                        0
11129                    } else {
11130                        ((k + r * 3 + g) % 3) as u8 // 0,1,2
11131                    };
11132                    cortiq_core::quant::q1t_pack(&mut c, k, code);
11133                }
11134                bytes.extend_from_slice(&c);
11135            }
11136        }
11137        // Per-row overlay: [u32 row_ptr[rows+1]] then [(u16 col, f16 val)] by
11138        // row (outliers are sorted by flat index → already grouped by row).
11139        let mut row_ptr = vec![0u32; rows + 1];
11140        for &(idx, _) in &outliers {
11141            row_ptr[idx as usize / cols + 1] += 1;
11142        }
11143        for r in 0..rows {
11144            row_ptr[r + 1] += row_ptr[r];
11145        }
11146        for &p in &row_ptr {
11147            bytes.extend_from_slice(&p.to_le_bytes());
11148        }
11149        for &(idx, v) in &outliers {
11150            bytes.extend_from_slice(&((idx as usize % cols) as u16).to_le_bytes());
11151            bytes.extend_from_slice(&f32_to_f16(v).to_le_bytes());
11152        }
11153
11154        let mut refw = vec![0f32; rows * cols];
11155        dequant_q1t(&bytes, rows, cols, &mut refw);
11156        // On-grid activations (±1, amax 1) so the int8 SDOT path reconstructs
11157        // x exactly and matches the f32 reference (same trick as the q1 test).
11158        let x: Vec<f32> = (0..cols)
11159            .map(|j| if j % 3 == 0 { 1.0 } else { -1.0 })
11160            .collect();
11161        let mut expect = vec![0f32; rows];
11162        for r in 0..rows {
11163            let mut a = 0.0f32;
11164            for j in 0..cols {
11165                a += refw[r * cols + j] * x[j];
11166            }
11167            expect[r] = a;
11168        }
11169        let tol = |e: f32| 1e-3 * e.abs().max(1e-3);
11170        let mut got = vec![0f32; rows];
11171        q1t_matvec(&bytes, &x, rows, cols, &mut got, None);
11172        for r in 0..rows {
11173            assert!(
11174                (got[r] - expect[r]).abs() < tol(expect[r]),
11175                "row {r}: {} vs {}",
11176                got[r],
11177                expect[r]
11178            );
11179        }
11180        // matmat (b=2, f32 decode path) must agree too.
11181        let x2: Vec<f32> = x.iter().chain(x.iter()).copied().collect();
11182        let mut gm = vec![0f32; 2 * rows];
11183        q1t_matmat(&bytes, &x2, 2, rows, cols, &mut gm, None);
11184        for r in 0..rows {
11185            assert!((gm[r] - expect[r]).abs() < tol(expect[r]));
11186            assert!((gm[rows + r] - expect[r]).abs() < tol(expect[r]));
11187        }
11188        // Fused pair (q1t_matvec2) must equal two single matvecs
11189        // bit-for-bit: same unpack, same group order, same f32
11190        // accumulation per stream. Distinct x2 exercises both lanes.
11191        let xb: Vec<f32> = (0..cols)
11192            .map(|j| if j % 5 == 0 { -1.0 } else { 1.0 })
11193            .collect();
11194        let (mut s1, mut s2) = (vec![0f32; rows], vec![0f32; rows]);
11195        q1t_matvec(&bytes, &x, rows, cols, &mut s1, None);
11196        q1t_matvec(&bytes, &xb, rows, cols, &mut s2, None);
11197        let (mut p1, mut p2) = (vec![0f32; rows], vec![0f32; rows]);
11198        q1t_matvec2(&bytes, &x, &xb, rows, cols, &mut p1, &mut p2, None);
11199        assert_eq!(p1, s1, "q1t pair lane 1 ≠ single matvec");
11200        assert_eq!(p2, s2, "q1t pair lane 2 ≠ single matvec");
11201    }
11202
11203    /// Pair == 2×matvec with an ODD group count (the kernel's tail
11204    /// group) and no overlay section.
11205    #[test]
11206    fn q1t_matvec2_odd_gpr_matches_singles() {
11207        use cortiq_core::quant::{Q1T_TILE, f32_to_f16, q1t_pack};
11208        let (rows, cols) = (5usize, 96usize); // gpr = 3 → paired + tail
11209        let gpr = cols / GROUP_SIZE;
11210        let mut bytes = Vec::with_capacity(rows * gpr * Q1T_TILE);
11211        for r in 0..rows {
11212            for g in 0..gpr {
11213                bytes.extend_from_slice(&f32_to_f16(0.1 + 0.05 * (r + g) as f32).to_le_bytes());
11214                let mut c = [0u8; 7];
11215                for k in 0..GROUP_SIZE {
11216                    q1t_pack(&mut c, k, ((k * 7 + r * 5 + g * 3) % 3) as u8);
11217                }
11218                bytes.extend_from_slice(&c);
11219            }
11220        }
11221        let x1: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.31).sin()).collect();
11222        let x2: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.17).cos()).collect();
11223        let (mut s1, mut s2) = (vec![0f32; rows], vec![0f32; rows]);
11224        q1t_matvec(&bytes, &x1, rows, cols, &mut s1, None);
11225        q1t_matvec(&bytes, &x2, rows, cols, &mut s2, None);
11226        let (mut p1, mut p2) = (vec![0f32; rows], vec![0f32; rows]);
11227        q1t_matvec2(&bytes, &x1, &x2, rows, cols, &mut p1, &mut p2, None);
11228        assert_eq!(p1, s1, "odd-gpr pair lane 1 ≠ single");
11229        assert_eq!(p2, s2, "odd-gpr pair lane 2 ≠ single");
11230    }
11231
11232    // Speed A/B: fused pair (one unpack, two streams) vs two single
11233    // matvecs. Single-threaded, FFN-sized, min-of paired in-process.
11234    //   cargo test -p cortiq-engine --release q1t_matvec2_speed -- --ignored --nocapture
11235    #[test]
11236    #[ignore]
11237    fn q1t_matvec2_speed() {
11238        use cortiq_core::quant::{Q1T_TILE, f32_to_f16, q1t_pack};
11239        use std::time::Instant;
11240        let (rows, cols) = (8192usize, 4096usize);
11241        let gpr = cols / GROUP_SIZE;
11242        let mut bytes = Vec::with_capacity(rows * gpr * Q1T_TILE);
11243        for r in 0..rows {
11244            for g in 0..gpr {
11245                let s = 0.1 + ((r + g) % 7) as f32 * 0.01;
11246                bytes.extend_from_slice(&f32_to_f16(s).to_le_bytes());
11247                let mut c = [0u8; 7];
11248                for k in 0..GROUP_SIZE {
11249                    q1t_pack(&mut c, k, ((k * 7 + r + g) % 3) as u8);
11250                }
11251                bytes.extend_from_slice(&c);
11252            }
11253        }
11254        let x1: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.31).sin()).collect();
11255        let x2: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.17).cos()).collect();
11256        let (mut s1, mut s2) = (vec![0f32; rows], vec![0f32; rows]);
11257        let (mut p1, mut p2) = (vec![0f32; rows], vec![0f32; rows]);
11258        // Warm both paths once.
11259        q1t_matvec(&bytes, &x1, rows, cols, &mut s1, None);
11260        q1t_matvec2(&bytes, &x1, &x2, rows, cols, &mut p1, &mut p2, None);
11261        let (mut t_pair, mut t_two) = (f64::MAX, f64::MAX);
11262        for _ in 0..8 {
11263            let t0 = Instant::now();
11264            q1t_matvec2(&bytes, &x1, &x2, rows, cols, &mut p1, &mut p2, None);
11265            t_pair = t_pair.min(t0.elapsed().as_secs_f64() * 1000.0);
11266            let t1 = Instant::now();
11267            q1t_matvec(&bytes, &x1, rows, cols, &mut s1, None);
11268            q1t_matvec(&bytes, &x2, rows, cols, &mut s2, None);
11269            t_two = t_two.min(t1.elapsed().as_secs_f64() * 1000.0);
11270        }
11271        assert_eq!(p1, s1);
11272        assert_eq!(p2, s2);
11273        println!("q1t pair {rows}x{cols}: fused {t_pair:.2} ms | two singles {t_two:.2} ms");
11274    }
11275
11276    // Speed A/B: the base-3-division decode (what the packing commit left in
11277    // place) vs the fused sign-LUT matvec. Both single-threaded, same bytes.
11278    //   cargo test -p cortiq-engine q1t_matvec_speed -- --ignored --nocapture
11279    #[test]
11280    #[ignore]
11281    fn q1t_matvec_speed() {
11282        use cortiq_core::quant::{Q1T_TILE, f32_to_f16, q1t_code, q1t_pack};
11283        use std::time::Instant;
11284        let (rows, cols) = (8192usize, 4096usize); // FFN-sized
11285        let gpr = cols / GROUP_SIZE;
11286        let mut bytes = Vec::with_capacity(rows * gpr * Q1T_TILE + 16);
11287        for r in 0..rows {
11288            for g in 0..gpr {
11289                let s = 0.1 + ((r + g) % 7) as f32 * 0.01;
11290                bytes.extend_from_slice(&f32_to_f16(s).to_le_bytes());
11291                let mut c = [0u8; 7];
11292                for k in 0..GROUP_SIZE {
11293                    q1t_pack(&mut c, k, ((k * 7 + r + g) % 3) as u8);
11294                }
11295                bytes.extend_from_slice(&c);
11296            }
11297        }
11298        let (n, stride) = (rows * cols, 40usize); // ~2.5% outliers, per-row overlay
11299        let mut row_ptr = vec![0u32; rows + 1];
11300        let mut idx = 0usize;
11301        while idx < n {
11302            row_ptr[idx / cols + 1] += 1;
11303            idx += stride;
11304        }
11305        for r in 0..rows {
11306            row_ptr[r + 1] += row_ptr[r];
11307        }
11308        for &p in &row_ptr {
11309            bytes.extend_from_slice(&p.to_le_bytes());
11310        }
11311        let mut idx = 0usize;
11312        while idx < n {
11313            bytes.extend_from_slice(&((idx % cols) as u16).to_le_bytes());
11314            bytes.extend_from_slice(&f32_to_f16((idx % 13) as f32 * 0.1 - 0.6).to_le_bytes());
11315            idx += stride;
11316        }
11317        // On-grid ±1 so the fast path's int8 SDOT is exact vs the f32 "slow"
11318        // reference (the A/B is a timing check; values must still agree).
11319        let x: Vec<f32> = (0..cols)
11320            .map(|j| if j % 3 == 0 { 1.0 } else { -1.0 })
11321            .collect();
11322        let (rp_off, ent_off, has_ov) = q1t_overlay(&bytes, rows * gpr * Q1T_TILE, rows);
11323
11324        // "before": base-3 division decode into a buffer, then dot.
11325        let slow = |out: &mut [f32]| {
11326            let mut buf = vec![0f32; cols];
11327            for r in 0..rows {
11328                for g in 0..gpr {
11329                    let off = (r * gpr + g) * Q1T_TILE;
11330                    let s = f16_to_f32(u16::from_le_bytes([bytes[off], bytes[off + 1]]));
11331                    let codes = &bytes[off + 2..off + Q1T_TILE];
11332                    for k in 0..GROUP_SIZE {
11333                        buf[g * GROUP_SIZE + k] = match q1t_code(codes, k) {
11334                            1 => s,
11335                            2 => -s,
11336                            _ => 0.0,
11337                        };
11338                    }
11339                }
11340                out[r] = q1t_row_outlier_correction(&bytes, r, rp_off, ent_off, has_ov, &x)
11341                    + (0..cols).map(|j| buf[j] * x[j]).sum::<f32>();
11342            }
11343        };
11344        let iters = 5;
11345        let mut a = vec![0f32; rows];
11346        slow(&mut a); // warm
11347        let t = Instant::now();
11348        for _ in 0..iters {
11349            slow(&mut a);
11350        }
11351        let slow_ms = t.elapsed().as_secs_f64() * 1e3 / iters as f64;
11352
11353        let mut b = vec![0f32; rows];
11354        q1t_matvec(&bytes, &x, rows, cols, &mut b, None); // warm
11355        let t = Instant::now();
11356        for _ in 0..iters {
11357            q1t_matvec(&bytes, &x, rows, cols, &mut b, None);
11358        }
11359        let fast_ms = t.elapsed().as_secs_f64() * 1e3 / iters as f64;
11360
11361        for r in 0..rows {
11362            assert!((a[r] - b[r]).abs() < 1e-2, "mismatch row {r}");
11363        }
11364        println!(
11365            "q1t matvec {rows}x{cols} (1 thread): div-decode {slow_ms:.2} ms  fused-LUT {fast_ms:.2} ms  => {:.2}x",
11366            slow_ms / fast_ms
11367        );
11368    }
11369}
11370
11371#[cfg(test)]
11372mod gemm_bench {
11373    /// `cargo test -p cortiq-engine --release q4tp_matmat_throughput -- --ignored --nocapture`
11374    /// Times the batched q4tp GEMM at the shapes the image DiT runs
11375    /// (b=296 tokens, 2304 -> 9216), on synthetic bytes: no model, no
11376    /// mmap, no thermal drift over minutes — a kernel change shows up
11377    /// here in seconds where a full render hides it in noise.
11378    ///
11379    /// On macOS add `CMF_ACCEL=0`: this shape is over the 500k-cell mark
11380    /// where the matmat hands off to Accelerate's dequant sgemm, and
11381    /// without the opt-out both rows below measure the AMX, not the
11382    /// kernel under test.
11383    #[test]
11384    #[ignore]
11385    fn q4tp_matmat_throughput() {
11386        // 296 is a prompt-encode batch; the image DiT runs 2085 at
11387        // 512x512, where the activation panel stops fitting L2 and the
11388        // loop's shape starts to matter more than its instructions.
11389        let b: usize = std::env::var("CMF_BENCH_B")
11390            .ok()
11391            .and_then(|v| v.parse().ok())
11392            .unwrap_or(296);
11393        let (rows, cols) = (9216usize, 2304usize);
11394        let (_, _, _) = (rows, cols, b);
11395        let total =
11396            cortiq_core::quant::expected_nbytes(cortiq_core::TensorDtype::Q4TiledP, &[rows, cols])
11397                .unwrap();
11398        // Random nibbles are fine, but the row params are f16 (lo, step)
11399        // of a geometric ladder: garbage there gives exp2 of a huge
11400        // exponent, the scales come back inf, and the whole bench times
11401        // NaN arithmetic instead of the kernel.
11402        let (params_off, codes_off, _) = cortiq_core::quant::q4tp_sections(rows, cols);
11403        let mut bytes: Vec<u8> = (0..total).map(|i| (i * 37 % 251) as u8).collect();
11404        let lo = cortiq_core::quant::f32_to_f16(-4.0);
11405        let step = cortiq_core::quant::f32_to_f16(0.1);
11406        for r in 0..rows {
11407            let o = params_off + r * 4;
11408            bytes[o..o + 2].copy_from_slice(&lo.to_le_bytes());
11409            bytes[o + 2..o + 4].copy_from_slice(&step.to_le_bytes());
11410        }
11411        let _ = codes_off;
11412        let xs: Vec<f32> = (0..b * cols)
11413            .map(|i| ((i % 97) as f32 - 48.0) / 48.0)
11414            .collect();
11415        let mut out = vec![0f32; b * rows];
11416        let pool = crate::pool::Pool::from_env();
11417        // A shared 48-core stand drifts ±25% run to run, which is wider
11418        // than any kernel change worth making. So: alternate the two
11419        // kernels inside one process and keep the BEST time for
11420        // each. Interleaving makes both see the same interference, and a
11421        // minimum is the one statistic another tenant cannot inflate.
11422        super::q4tp_matmat(&bytes, &xs, b, rows, cols, &mut out, pool.as_deref());
11423        let reps: usize = std::env::var("CMF_BENCH_REPS")
11424            .ok()
11425            .and_then(|v| v.parse().ok())
11426            .unwrap_or(10);
11427        let mut best = [f64::MAX; 2];
11428        let mut sums = [0f32; 2];
11429        for _ in 0..reps {
11430            for (k, w) in [(0usize, 1u8), (1usize, 2u8)] {
11431                super::Q4TP_ALT.store(w, std::sync::atomic::Ordering::Relaxed);
11432                let t = std::time::Instant::now();
11433                super::q4tp_matmat(&bytes, &xs, b, rows, cols, &mut out, pool.as_deref());
11434                best[k] = best[k].min(t.elapsed().as_secs_f64());
11435                sums[k] = out.iter().take(64).sum::<f32>();
11436            }
11437        }
11438        let flops = 2.0 * b as f64 * rows as f64 * cols as f64;
11439        for (k, name) in ["previous", "tuned   "].iter().enumerate() {
11440            println!(
11441                "q4tp matmat {rows}x{cols} b={b} {name}: {:.1} ms  {:.1} GFLOP/s  (checksum {:.3})",
11442                best[k] * 1e3,
11443                flops / best[k] / 1e9,
11444                sums[k]
11445            );
11446        }
11447        assert!(
11448            (sums[0] - sums[1]).abs() < 1e-2,
11449            "the tuned kernel changed the result: {} vs {}",
11450            sums[0],
11451            sums[1]
11452        );
11453    }
11454
11455    /// The blocked kernel must agree with the per-column path exactly —
11456    /// same weights, same activation split, only a different instruction
11457    /// mix. Shapes are chosen to hit the awkward cases: a column count
11458    /// that leaves an odd group (the 512-bit kernel does two at a time),
11459    /// and a batch that does not divide by four.
11460    #[test]
11461    fn q4tp_matmat_blocked_matches_scalar() {
11462        use std::sync::atomic::Ordering::Relaxed;
11463        // The last shape carries the image DiT's column count — 2304, so
11464        // 72 groups of accumulation, which is where a reordered sum can
11465        // actually drift — and runs through the thread pool, since the
11466        // blocked path splits rows across workers. Its row count stays
11467        // under 500k cells on purpose: above that, macOS diverts the whole
11468        // matmat to the Accelerate/AMX dequant sgemm and neither kernel
11469        // here would run.
11470        for &(rows, cols, b) in &[
11471            (64usize, 128usize, 7usize),
11472            (33, 96, 4),
11473            (16, 256, 9),
11474            (192, 2304, 37),
11475        ] {
11476            let total = cortiq_core::quant::expected_nbytes(
11477                cortiq_core::TensorDtype::Q4TiledP,
11478                &[rows, cols],
11479            )
11480            .unwrap();
11481            let (params_off, _, _) = cortiq_core::quant::q4tp_sections(rows, cols);
11482            let mut bytes: Vec<u8> = (0..total).map(|i| (i * 61 % 251) as u8).collect();
11483            let lo = cortiq_core::quant::f32_to_f16(-4.0);
11484            let step = cortiq_core::quant::f32_to_f16(0.1);
11485            for r in 0..rows {
11486                let o = params_off + r * 4;
11487                bytes[o..o + 2].copy_from_slice(&lo.to_le_bytes());
11488                bytes[o + 2..o + 4].copy_from_slice(&step.to_le_bytes());
11489            }
11490            let xs: Vec<f32> = (0..b * cols)
11491                .map(|i| ((i % 89) as f32 - 44.0) / 44.0)
11492                .collect();
11493            let mut got = vec![0f32; b * rows];
11494            let mut want = vec![0f32; b * rows];
11495            let gpr = cols / 32;
11496            let view = super::Q4tpView::new(&bytes, rows, cols);
11497            let pool = crate::pool::Pool::from_env();
11498            super::Q4TP_ALT.store(2, Relaxed);
11499            super::q4tp_matmat(&bytes, &xs, b, rows, cols, &mut got, pool.as_deref());
11500            super::Q4TP_ALT.store(1, Relaxed);
11501            super::q4tp_matmat(&bytes, &xs, b, rows, cols, &mut want, pool.as_deref());
11502            super::Q4TP_ALT.store(0, Relaxed);
11503            // Measured against the output's scale, not cell by cell: a
11504            // dot product of 2304 terms lands near zero wherever the row
11505            // and the activation nearly cancel, and there a per-cell
11506            // ratio reports 1e-3 for an absolute error of 5e-6 — f32's
11507            // own rounding, reordered. What must stay small is the error
11508            // relative to what the layer actually outputs.
11509            let scale = want.iter().fold(0f32, |m, v| m.max(v.abs())).max(1e-6);
11510            let (mut worst, mut at) = (0f32, 0usize);
11511            for (i, (g, w)) in got.iter().zip(&want).enumerate() {
11512                if (g - w).abs() > worst {
11513                    worst = (g - w).abs();
11514                    at = i;
11515                }
11516            }
11517            assert!(
11518                worst <= 1e-4 * scale,
11519                "{rows}x{cols} b={b}: blocked and scalar disagree by {worst:.3e} \
11520                 (scale {scale:.3e}) at cell {at}: {} vs {}",
11521                got[at],
11522                want[at]
11523            );
11524
11525            // "Same speed, no quality loss" is a claim about which answer
11526            // is RIGHT, not about which two agree. Both paths sum the same
11527            // 2304 products in different orders, so f64 decides: the
11528            // blocked kernel keeps sixteen partial sums and folds them at
11529            // the end, which is a shallower addition tree than the
11530            // per-column path's running scalar, and it must not be worse.
11531            let (mut e_blocked, mut e_scalar) = (0f64, 0f64);
11532            for bi in 0..b {
11533                let act = super::split_act(&xs[bi * cols..(bi + 1) * cols]);
11534                for r in 0..rows {
11535                    let mut sc = vec![0f32; gpr];
11536                    view.scales_into(r, gpr, &mut sc);
11537                    let mut exact = 0f64;
11538                    for j in 0..cols {
11539                        let (w, sq) = super::q4tp_outlier(view.nib, r, gpr, j, &sc);
11540                        exact += w as f64 * sq as f64 * act.xq[j] as f64;
11541                    }
11542                    exact *= act.sx as f64;
11543                    for &(j, xv) in &act.outliers {
11544                        let (w, sq) = super::q4tp_outlier(view.nib, r, gpr, j, &sc);
11545                        exact += w as f64 * sq as f64 * xv as f64;
11546                    }
11547                    let i = bi * rows + r;
11548                    e_blocked = e_blocked.max((got[i] as f64 - exact).abs());
11549                    e_scalar = e_scalar.max((want[i] as f64 - exact).abs());
11550                }
11551            }
11552            println!(
11553                "{rows}x{cols} b={b}: worst error vs f64 — blocked {e_blocked:.3e}, \
11554                 per-column {e_scalar:.3e}"
11555            );
11556            // An absolute bar, not a race between the two: at these
11557            // magnitudes both sit in f32's last bits, and on a small shape
11558            // whichever one happens to round the unluckiest cell "wins" by
11559            // a factor the next seed reverses.
11560            assert!(
11561                e_blocked <= 1e-5 * scale as f64 && e_scalar <= 1e-5 * scale as f64,
11562                "{rows}x{cols} b={b}: error against f64 too large — blocked \
11563                 {e_blocked:.3e}, per-column {e_scalar:.3e}, scale {scale:.3e}"
11564            );
11565        }
11566    }
11567
11568    /// The q4t twin of the throughput bench, same shape and rules, so the
11569    /// two quantisations' batch kernels can be read against each other.
11570    /// `cargo test -p cortiq-engine --release q4t_matmat_throughput -- --ignored --nocapture`
11571    #[test]
11572    #[ignore]
11573    fn q4t_matmat_throughput() {
11574        let (rows, cols, b) = (9216usize, 2304usize, 296usize);
11575        let total =
11576            cortiq_core::quant::expected_nbytes(cortiq_core::TensorDtype::Q4Tiled, &[rows, cols])
11577                .unwrap();
11578        // q4t carries a per-group f16 scale in the tile's first two bytes;
11579        // random bytes there decode to inf and the bench would time NaNs.
11580        let mut bytes: Vec<u8> = (0..total).map(|i| (i * 37 % 251) as u8).collect();
11581        let sc = cortiq_core::quant::f32_to_f16(0.02);
11582        for t in bytes.chunks_mut(super::Q4_TILE) {
11583            t[..2].copy_from_slice(&sc.to_le_bytes());
11584        }
11585        let xs: Vec<f32> = (0..b * cols)
11586            .map(|i| ((i % 97) as f32 - 48.0) / 48.0)
11587            .collect();
11588        let mut out = vec![0f32; b * rows];
11589        let pool = crate::pool::Pool::from_env();
11590        super::q4t_matmat(&bytes, &xs, b, rows, cols, &mut out, pool.as_deref());
11591        let reps: usize = std::env::var("CMF_BENCH_REPS")
11592            .ok()
11593            .and_then(|v| v.parse().ok())
11594            .unwrap_or(10);
11595        let mut best = f64::MAX;
11596        for _ in 0..reps {
11597            let t = std::time::Instant::now();
11598            super::q4t_matmat(&bytes, &xs, b, rows, cols, &mut out, pool.as_deref());
11599            best = best.min(t.elapsed().as_secs_f64());
11600        }
11601        let flops = 2.0 * b as f64 * rows as f64 * cols as f64;
11602        println!(
11603            "q4t matmat {rows}x{cols} b={b}: {:.1} ms  {:.1} GFLOP/s  (checksum {:.3})",
11604            best * 1e3,
11605            flops / best / 1e9,
11606            out.iter().take(64).sum::<f32>()
11607        );
11608    }
11609}