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    /// The device GEMM this tensor would take, run once on the caller's
1770    /// data — the startup parity probe's arm, and the one place that knows
1771    /// which entry point each codec has.
1772    ///
1773    /// It exists because the probe used to look for a `q4tp` weight by
1774    /// name AND dtype, and a container packed any other way was declared
1775    /// "host path" for the whole render even though its codec had a device
1776    /// GEMM of its own. A gate that only recognizes one codec is a gate
1777    /// that silently downgrades every other one.
1778    pub fn device_matmat(&self, xs: &[f32], b: usize, out: &mut [f32]) -> bool {
1779        let (rows, cols) = (self.rows(), self.cols());
1780        let Self::Mapped { model, idx, dtype, row_scale, col_field, .. } = self else {
1781            return false;
1782        };
1783        match *dtype {
1784            TensorDtype::Q4TiledP => {
1785                crate::gpu::q4tp_matmat(model, *idx, xs, b, rows, cols, out)
1786            }
1787            // The two-field codec folds its column field into the
1788            // activation, which leaves a plain per-row int8 GEMM — the
1789            // same kernel `q8_row` uses, on both backends.
1790            TensorDtype::Q8Row | TensorDtype::Q8_2f => {
1791                let flat: Vec<f32> = (0..b)
1792                    .flat_map(|bi| {
1793                        prescale(&xs[bi * cols..(bi + 1) * cols], col_field, *dtype).into_owned()
1794                    })
1795                    .collect();
1796                crate::gpu::q8_matmat(model, *idx, row_scale, &flat, b, rows, cols, out)
1797            }
1798            _ => false,
1799        }
1800    }
1801
1802    /// Multi-matrix job (roadmap §3 P0): N tensors sharing one input
1803    /// run under a SINGLE pool dispatch — QKV or gate+up cost one
1804    /// barrier instead of N. Per-row math is the exact same kernel as
1805    /// `matvec` (bit-identical outputs); only the dispatch is fused.
1806    /// Falls back to N sequential matvecs when the set is not a uniform
1807    /// q8-family/F32 group or there is no pool.
1808    pub fn matvec_many<const N: usize>(
1809        ts: [&QTensor; N],
1810        x: &[f32],
1811        mut outs: [&mut [f32]; N],
1812        pool: Option<&Pool>,
1813    ) {
1814        let total_rows: usize = ts.iter().map(|t| t.rows()).sum();
1815        let uniform_q8 = ts.iter().all(|t| {
1816            matches!(
1817                t,
1818                Self::Mapped {
1819                    dtype: TensorDtype::Q8Row | TensorDtype::Q8_2f,
1820                    ..
1821                }
1822            )
1823        });
1824        let uniform_f32 = ts.iter().all(|t| matches!(t, Self::F32 { .. }));
1825        let uniform_q4 = ts.iter().all(|t| {
1826            matches!(
1827                t,
1828                Self::Mapped {
1829                    dtype: TensorDtype::Q4Block,
1830                    ..
1831                }
1832            )
1833        });
1834        let uniform_vbit = ts.iter().all(|t| {
1835            matches!(
1836                t,
1837                Self::Mapped {
1838                    dtype: TensorDtype::Vbit | TensorDtype::VbitRo,
1839                    ..
1840                }
1841            )
1842        });
1843        let uniform_q1 = ts.iter().all(|t| {
1844            matches!(
1845                t,
1846                Self::Mapped {
1847                    dtype: TensorDtype::Q1,
1848                    ..
1849                }
1850            )
1851        });
1852        let uniform_q1t = ts.iter().all(|t| {
1853            matches!(
1854                t,
1855                Self::Mapped {
1856                    dtype: TensorDtype::Q1T,
1857                    ..
1858                }
1859            )
1860        });
1861        // q4tp is the skeleton dtype of the big MoE files, and without an arm
1862        // here every projection that shares an input paid its own pool
1863        // barrier: DeepSeek-V4's attention step alone hands this function
1864        // wq_a, wkv and both compressors' pairs off the same hidden state.
1865        let uniform_q4tp = ts.iter().all(|t| {
1866            matches!(
1867                t,
1868                Self::Mapped {
1869                    dtype: TensorDtype::Q4TiledP,
1870                    ..
1871                }
1872            )
1873        }) && ts
1874            .iter()
1875            .all(|t| t.cols() == ts[0].cols() && t.cols() % GROUP_SIZE == 0);
1876        let Some(pool) = pool else {
1877            for (t, o) in ts.iter().zip(outs.iter_mut()) {
1878                t.matvec(x, o, None);
1879            }
1880            return;
1881        };
1882        if total_rows < 256
1883            || !(uniform_q8
1884                || uniform_f32
1885                || uniform_q4
1886                || uniform_vbit
1887                || uniform_q1
1888                || uniform_q1t
1889                || uniform_q4tp)
1890        {
1891            for (t, o) in ts.iter().zip(outs.iter_mut()) {
1892                t.matvec(x, o, Some(pool));
1893            }
1894            return;
1895        }
1896
1897        if uniform_q4tp {
1898            // Every tensor's rows laid end to end in one virtual row space,
1899            // so the whole set is ONE dispatch. The per-row body is the
1900            // `q4tp_matvec` arm verbatim — same activation split, same
1901            // accumulation order — so the outputs are bit-identical to the
1902            // sequential calls this replaces.
1903            let cols = ts[0].cols();
1904            let gpr = cols / GROUP_SIZE;
1905            let views: Vec<Q4tpView> = ts
1906                .iter()
1907                .map(|t| Q4tpView::new(t.quant_bytes(), t.rows(), cols))
1908                .collect();
1909            let rows_of: Vec<usize> = ts.iter().map(|t| t.rows()).collect();
1910            let outs_addr: [SendMut; N] = std::array::from_fn(|i| SendMut(outs[i].as_mut_ptr()));
1911            // flat index -> (which tensor, which of its rows)
1912            let locate = |flat: usize| -> (usize, usize) {
1913                let mut acc = 0;
1914                for (i, &r) in rows_of.iter().enumerate() {
1915                    if flat < acc + r {
1916                        return (i, flat - acc);
1917                    }
1918                    acc += r;
1919                }
1920                (rows_of.len() - 1, 0)
1921            };
1922            let (views, outs_addr) = (&views, &outs_addr);
1923            if a8w8_enabled() {
1924                let act = split_act(x);
1925                let act = &act;
1926                let run = |start: usize, end: usize| {
1927                    let mut sc = vec![0f32; gpr];
1928                    for flat in start..end {
1929                        let (t, r) = locate(flat);
1930                        let v = &views[t];
1931                        v.scales_into(r, gpr, &mut sc);
1932                        let mut acc = dot_q4tp_row_i8(v.nib, r, gpr, &act.xq, &sc) * act.sx;
1933                        for &(j, xv) in &act.outliers {
1934                            let (w, s) = q4tp_outlier(v.nib, r, gpr, j, &sc);
1935                            acc += w * s * xv;
1936                        }
1937                        // SAFETY: one worker owns each (tensor, row) pair.
1938                        unsafe { *outs_addr[t].at(r) = acc };
1939                    }
1940                };
1941                pool.run_rows(total_rows, &run);
1942            } else {
1943                let run = |start: usize, end: usize| {
1944                    let mut sc = vec![0f32; gpr];
1945                    for flat in start..end {
1946                        let (t, r) = locate(flat);
1947                        let v = &views[t];
1948                        v.scales_into(r, gpr, &mut sc);
1949                        // SAFETY: one worker owns each (tensor, row) pair.
1950                        unsafe { *outs_addr[t].at(r) = q4tp_row_exact(v.nib, r, gpr, x, &sc) };
1951                    }
1952                };
1953                pool.run_rows(total_rows, &run);
1954            }
1955            return;
1956        }
1957
1958        if uniform_q1 {
1959            // One shared activation split + group sums (q1 has no col
1960            // field; the same input feeds every tensor).
1961            let outs_addr: [SendMut; N] = std::array::from_fn(|i| SendMut(outs[i].as_mut_ptr()));
1962            if a8w8_enabled() {
1963                let act = split_act(x);
1964                let gsum = q1_group_sums(&act.xq, ts[0].cols() / GROUP_SIZE);
1965                let (act, gsum) = (&act, &gsum);
1966                let closures: [_; N] = std::array::from_fn(|i| {
1967                    let (bytes, gpr, out) =
1968                        (ts[i].quant_bytes(), ts[i].cols() / GROUP_SIZE, outs_addr[i]);
1969                    move |s: usize, e: usize| q1_range_a8w8(bytes, gpr, act, gsum, out, s, e)
1970                });
1971                let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
1972                    std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
1973                pool.run_many(&parts);
1974            } else {
1975                let closures: [_; N] = std::array::from_fn(|i| {
1976                    let (bytes, gpr, out) =
1977                        (ts[i].quant_bytes(), ts[i].cols() / GROUP_SIZE, outs_addr[i]);
1978                    move |s: usize, e: usize| q1_range_f32(bytes, gpr, x, out, s, e)
1979                });
1980                let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
1981                    std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
1982                pool.run_many(&parts);
1983            }
1984            return;
1985        }
1986
1987        if uniform_q1t {
1988            // Q1T batched: one shared activation split + overlay decode,
1989            // all tensors' rows in ONE pool dispatch (saves N−1 dispatches
1990            // and N−1 redundant split_act calls per layer).
1991            let outs_addr: [SendMut; N] = std::array::from_fn(|i| SendMut(outs[i].as_mut_ptr()));
1992            const TILE: usize = cortiq_core::quant::Q1T_TILE;
1993            if a8w8_enabled() {
1994                let act = split_act(x);
1995                let act = &act;
1996                let x_ref = x;
1997                let closures: [_; N] = std::array::from_fn(|i| {
1998                    let bytes = ts[i].quant_bytes();
1999                    let (rows, cols) = (ts[i].rows(), ts[i].cols());
2000                    let gpr = cols / GROUP_SIZE;
2001                    let (rp_off, ent_off, has_ov) = q1t_overlay(bytes, rows * gpr * TILE, rows);
2002                    let out = outs_addr[i];
2003                    move |s: usize, e: usize| {
2004                        q1t_range_a8w8(bytes, gpr, rp_off, ent_off, has_ov, act, x_ref, out, s, e)
2005                    }
2006                });
2007                let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
2008                    std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
2009                pool.run_many(&parts);
2010            } else {
2011                let x_ref = x;
2012                let closures: [_; N] = std::array::from_fn(|i| {
2013                    let bytes = ts[i].quant_bytes();
2014                    let (rows, cols) = (ts[i].rows(), ts[i].cols());
2015                    let gpr = cols / GROUP_SIZE;
2016                    let (rp_off, ent_off, has_ov) = q1t_overlay(bytes, rows * gpr * TILE, rows);
2017                    let out = outs_addr[i];
2018                    move |s: usize, e: usize| {
2019                        q1t_range_f32_batch(bytes, gpr, rp_off, ent_off, has_ov, x_ref, out, s, e)
2020                    }
2021                });
2022                let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
2023                    std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
2024                pool.run_many(&parts);
2025            }
2026            return;
2027        }
2028
2029        if uniform_q4 || uniform_vbit {
2030            let outs_addr: [SendMut; N] = std::array::from_fn(|i| SendMut(outs[i].as_mut_ptr()));
2031            // q4/vbit share one activation split — no per-tensor col field.
2032            if a8w8_enabled() {
2033                let act = split_act(x);
2034                let act = &act;
2035                if uniform_q4 {
2036                    let closures: [_; N] = std::array::from_fn(|i| {
2037                        let (packed, scales) =
2038                            q4_split(ts[i].quant_bytes(), ts[i].rows(), ts[i].cols());
2039                        let (gpr, cols, out) =
2040                            (ts[i].cols() / GROUP_SIZE, ts[i].cols(), outs_addr[i]);
2041                        move |s: usize, e: usize| {
2042                            q4_range_a8w8(packed, scales, gpr, cols, act, out, s, e)
2043                        }
2044                    });
2045                    let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
2046                        std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
2047                    pool.run_many(&parts);
2048                } else {
2049                    let closures: [_; N] = std::array::from_fn(|i| {
2050                        let Self::Mapped { vbit_offsets, .. } = ts[i] else {
2051                            unreachable!()
2052                        };
2053                        let (bytes, rows, cols, out) = (
2054                            ts[i].quant_bytes(),
2055                            ts[i].rows(),
2056                            ts[i].cols(),
2057                            outs_addr[i],
2058                        );
2059                        move |s: usize, e: usize| {
2060                            vbit_range_a8w8(bytes, vbit_offsets, x, act, rows, cols, out, s, e)
2061                        }
2062                    });
2063                    let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
2064                        std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
2065                    pool.run_many(&parts);
2066                }
2067                return;
2068            }
2069            if uniform_q4 {
2070                let closures: [_; N] = std::array::from_fn(|i| {
2071                    let (packed, scales) =
2072                        q4_split(ts[i].quant_bytes(), ts[i].rows(), ts[i].cols());
2073                    let (gpr, out) = (ts[i].cols() / GROUP_SIZE, outs_addr[i]);
2074                    move |s: usize, e: usize| q4_range_f32(packed, scales, gpr, x, out, s, e)
2075                });
2076                let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
2077                    std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
2078                pool.run_many(&parts);
2079            } else {
2080                let closures: [_; N] = std::array::from_fn(|i| {
2081                    let Self::Mapped { vbit_offsets, .. } = ts[i] else {
2082                        unreachable!()
2083                    };
2084                    let (bytes, rows, cols, out) = (
2085                        ts[i].quant_bytes(),
2086                        ts[i].rows(),
2087                        ts[i].cols(),
2088                        outs_addr[i],
2089                    );
2090                    move |s: usize, e: usize| {
2091                        vbit_range_f32(bytes, vbit_offsets, x, rows, cols, out, s, e)
2092                    }
2093                });
2094                let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
2095                    std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
2096                pool.run_many(&parts);
2097            }
2098            return;
2099        }
2100
2101        if uniform_f32 {
2102            let outs_addr: [SendMut; N] = std::array::from_fn(|i| SendMut(outs[i].as_mut_ptr()));
2103            let closures: [_; N] = std::array::from_fn(|i| {
2104                let Self::F32 { data, cols, .. } = ts[i] else {
2105                    unreachable!()
2106                };
2107                let out = outs_addr[i];
2108                move |start: usize, end: usize| {
2109                    for o in start..end {
2110                        let row = &data[o * cols..(o + 1) * cols];
2111                        let mut sum = 0.0f32;
2112                        for j in 0..*cols {
2113                            sum += row[j] * x[j];
2114                        }
2115                        // SAFETY: disjoint (tensor, row) cells per worker.
2116                        unsafe { *out.at(o) = sum };
2117                    }
2118                }
2119            });
2120            let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
2121                std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
2122            pool.run_many(&parts);
2123            return;
2124        }
2125
2126        // Uniform q8-family: per-tensor prescale (q8_2f col fields
2127        // differ per tensor) + the shared range kernels.
2128        struct Ctx<'a> {
2129            bytes: &'a [u8],
2130            #[cfg_attr(not(target_arch = "aarch64"), allow(dead_code))]
2131            rep: &'a [u8],
2132            row_scale: &'a [f32],
2133            cols: usize,
2134            xs: std::borrow::Cow<'a, [f32]>,
2135        }
2136        let ctxs: [Ctx<'_>; N] = std::array::from_fn(|i| {
2137            let Self::Mapped {
2138                dtype,
2139                cols,
2140                row_scale,
2141                col_field,
2142                repack,
2143                ..
2144            } = ts[i]
2145            else {
2146                unreachable!()
2147            };
2148            Ctx {
2149                bytes: ts[i].quant_bytes(),
2150                rep: repack,
2151                row_scale,
2152                cols: *cols,
2153                xs: prescale(x, col_field, *dtype),
2154            }
2155        });
2156        let outs_addr: [SendMut; N] = std::array::from_fn(|i| SendMut(outs[i].as_mut_ptr()));
2157        #[cfg(target_arch = "aarch64")]
2158        if sdot_enabled() {
2159            let acts: [SplitAct; N] = std::array::from_fn(|i| split_act(&ctxs[i].xs));
2160            let closures: [_; N] = std::array::from_fn(|i| {
2161                let (c, act, out) = (&ctxs[i], &acts[i], outs_addr[i]);
2162                move |start: usize, end: usize| {
2163                    q8_range_sdot(c.bytes, c.rep, c.row_scale, act, c.cols, out, start, end)
2164                }
2165            });
2166            let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
2167                std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
2168            pool.run_many(&parts);
2169            return;
2170        }
2171        #[cfg(target_arch = "x86_64")]
2172        if avx2_a8w8_enabled() {
2173            let acts: [SplitAct; N] = std::array::from_fn(|i| split_act(&ctxs[i].xs));
2174            let closures: [_; N] = std::array::from_fn(|i| {
2175                let (c, act, out) = (&ctxs[i], &acts[i], outs_addr[i]);
2176                move |start: usize, end: usize| {
2177                    q8_range_avx2(c.bytes, c.row_scale, act, c.cols, out, start, end)
2178                }
2179            });
2180            let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
2181                std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
2182            pool.run_many(&parts);
2183            return;
2184        }
2185        let closures: [_; N] = std::array::from_fn(|i| {
2186            let (c, out) = (&ctxs[i], outs_addr[i]);
2187            move |start: usize, end: usize| {
2188                q8_range_f32(c.bytes, c.row_scale, &c.xs, c.cols, out, start, end)
2189            }
2190        });
2191        let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
2192            std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
2193        pool.run_many(&parts);
2194    }
2195}
2196
2197impl QTensor {
2198    /// Pair-input multi-matrix job: N tensors × 2 shared inputs under a
2199    /// single pool dispatch — the MTP/pair decode path publishes one job
2200    /// for Q/K/V (and one for gate+up) instead of one per tensor.
2201    /// Per-row math is exactly `matvec2`'s kernels; bit-identical.
2202    #[allow(clippy::needless_range_loop)]
2203    pub fn matvec2_many<const N: usize>(
2204        ts: [&QTensor; N],
2205        x1: &[f32],
2206        x2: &[f32],
2207        mut o1s: [&mut [f32]; N],
2208        mut o2s: [&mut [f32]; N],
2209        pool: Option<&Pool>,
2210    ) {
2211        let total_rows: usize = ts.iter().map(|t| t.rows()).sum();
2212        let uniform_q8 = ts.iter().all(|t| {
2213            matches!(
2214                t,
2215                Self::Mapped {
2216                    dtype: TensorDtype::Q8Row | TensorDtype::Q8_2f,
2217                    ..
2218                }
2219            )
2220        });
2221        let uniform_f32 = ts.iter().all(|t| matches!(t, Self::F32 { .. }));
2222        let uniform_q4 = ts.iter().all(|t| {
2223            matches!(
2224                t,
2225                Self::Mapped {
2226                    dtype: TensorDtype::Q4Block,
2227                    ..
2228                }
2229            )
2230        });
2231        let uniform_vbit = ts.iter().all(|t| {
2232            matches!(
2233                t,
2234                Self::Mapped {
2235                    dtype: TensorDtype::Vbit | TensorDtype::VbitRo,
2236                    ..
2237                }
2238            )
2239        });
2240        let fusable = pool.is_some()
2241            && total_rows >= 256
2242            && (uniform_q8 || uniform_f32 || uniform_q4 || uniform_vbit);
2243        if !fusable {
2244            for i in 0..N {
2245                ts[i].matvec2(x1, x2, o1s[i], o2s[i], pool);
2246            }
2247            return;
2248        }
2249        let pool = pool.unwrap();
2250
2251        if uniform_q4 || uniform_vbit {
2252            let p1: [SendMut; N] = std::array::from_fn(|i| SendMut(o1s[i].as_mut_ptr()));
2253            let p2: [SendMut; N] = std::array::from_fn(|i| SendMut(o2s[i].as_mut_ptr()));
2254            // q4/vbit share activation splits — no per-tensor col field.
2255            if a8w8_enabled() {
2256                let a1 = split_act(x1);
2257                let a2 = split_act(x2);
2258                let (a1, a2) = (&a1, &a2);
2259                if uniform_q4 {
2260                    let closures: [_; N] = std::array::from_fn(|i| {
2261                        let (packed, scales) =
2262                            q4_split(ts[i].quant_bytes(), ts[i].rows(), ts[i].cols());
2263                        let (gpr, cols, o1, o2) =
2264                            (ts[i].cols() / GROUP_SIZE, ts[i].cols(), p1[i], p2[i]);
2265                        move |s: usize, e: usize| {
2266                            q4_range2_a8w8(packed, scales, gpr, cols, a1, a2, o1, o2, s, e)
2267                        }
2268                    });
2269                    let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
2270                        std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
2271                    pool.run_many(&parts);
2272                } else {
2273                    let closures: [_; N] = std::array::from_fn(|i| {
2274                        let Self::Mapped { vbit_offsets, .. } = ts[i] else {
2275                            unreachable!()
2276                        };
2277                        let (bytes, rows, cols, o1, o2) = (
2278                            ts[i].quant_bytes(),
2279                            ts[i].rows(),
2280                            ts[i].cols(),
2281                            p1[i],
2282                            p2[i],
2283                        );
2284                        move |s: usize, e: usize| {
2285                            vbit_range2_a8w8(
2286                                bytes,
2287                                vbit_offsets,
2288                                x1,
2289                                x2,
2290                                a1,
2291                                a2,
2292                                rows,
2293                                cols,
2294                                o1,
2295                                o2,
2296                                s,
2297                                e,
2298                            )
2299                        }
2300                    });
2301                    let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
2302                        std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
2303                    pool.run_many(&parts);
2304                }
2305                return;
2306            }
2307            if uniform_q4 {
2308                let closures: [_; N] = std::array::from_fn(|i| {
2309                    let (packed, scales) =
2310                        q4_split(ts[i].quant_bytes(), ts[i].rows(), ts[i].cols());
2311                    let (gpr, o1, o2) = (ts[i].cols() / GROUP_SIZE, p1[i], p2[i]);
2312                    move |s: usize, e: usize| {
2313                        q4_range2_f32(packed, scales, gpr, x1, x2, o1, o2, s, e)
2314                    }
2315                });
2316                let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
2317                    std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
2318                pool.run_many(&parts);
2319            } else {
2320                let closures: [_; N] = std::array::from_fn(|i| {
2321                    let Self::Mapped { vbit_offsets, .. } = ts[i] else {
2322                        unreachable!()
2323                    };
2324                    let (bytes, rows, cols, o1, o2) = (
2325                        ts[i].quant_bytes(),
2326                        ts[i].rows(),
2327                        ts[i].cols(),
2328                        p1[i],
2329                        p2[i],
2330                    );
2331                    move |s: usize, e: usize| {
2332                        vbit_range2_f32(bytes, vbit_offsets, x1, x2, rows, cols, o1, o2, s, e)
2333                    }
2334                });
2335                let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
2336                    std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
2337                pool.run_many(&parts);
2338            }
2339            return;
2340        }
2341
2342        if uniform_f32 {
2343            let p1: [SendMut; N] = std::array::from_fn(|i| SendMut(o1s[i].as_mut_ptr()));
2344            let p2: [SendMut; N] = std::array::from_fn(|i| SendMut(o2s[i].as_mut_ptr()));
2345            let closures: [_; N] = std::array::from_fn(|i| {
2346                let Self::F32 { data, cols, .. } = ts[i] else {
2347                    unreachable!()
2348                };
2349                let (o1, o2) = (p1[i], p2[i]);
2350                move |start: usize, end: usize| {
2351                    for o in start..end {
2352                        let row = &data[o * cols..(o + 1) * cols];
2353                        let (mut s1, mut s2) = (0.0f32, 0.0f32);
2354                        for j in 0..*cols {
2355                            s1 += row[j] * x1[j];
2356                            s2 += row[j] * x2[j];
2357                        }
2358                        // SAFETY: disjoint (tensor, row) cells per worker.
2359                        unsafe {
2360                            *o1.at(o) = s1;
2361                            *o2.at(o) = s2;
2362                        }
2363                    }
2364                }
2365            });
2366            let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
2367                std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
2368            pool.run_many(&parts);
2369            return;
2370        }
2371
2372        struct Ctx<'a> {
2373            bytes: &'a [u8],
2374            row_scale: &'a [f32],
2375            cols: usize,
2376            xs1: std::borrow::Cow<'a, [f32]>,
2377            xs2: std::borrow::Cow<'a, [f32]>,
2378        }
2379        let ctxs: [Ctx<'_>; N] = std::array::from_fn(|i| {
2380            let Self::Mapped {
2381                dtype,
2382                cols,
2383                row_scale,
2384                col_field,
2385                ..
2386            } = ts[i]
2387            else {
2388                unreachable!()
2389            };
2390            Ctx {
2391                bytes: ts[i].quant_bytes(),
2392                row_scale,
2393                cols: *cols,
2394                xs1: prescale(x1, col_field, *dtype),
2395                xs2: prescale(x2, col_field, *dtype),
2396            }
2397        });
2398        let p1: [SendMut; N] = std::array::from_fn(|i| SendMut(o1s[i].as_mut_ptr()));
2399        let p2: [SendMut; N] = std::array::from_fn(|i| SendMut(o2s[i].as_mut_ptr()));
2400        #[cfg(target_arch = "aarch64")]
2401        if sdot_enabled() {
2402            let acts: [(SplitAct, SplitAct); N] =
2403                std::array::from_fn(|i| (split_act(&ctxs[i].xs1), split_act(&ctxs[i].xs2)));
2404            let closures: [_; N] = std::array::from_fn(|i| {
2405                let (c, a, o1, o2) = (&ctxs[i], &acts[i], p1[i], p2[i]);
2406                move |start: usize, end: usize| {
2407                    q8_range2_sdot(c.bytes, c.row_scale, &a.0, &a.1, c.cols, o1, o2, start, end)
2408                }
2409            });
2410            let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
2411                std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
2412            pool.run_many(&parts);
2413            return;
2414        }
2415        #[cfg(target_arch = "x86_64")]
2416        if avx2_a8w8_enabled() {
2417            let acts: [(SplitAct, SplitAct); N] =
2418                std::array::from_fn(|i| (split_act(&ctxs[i].xs1), split_act(&ctxs[i].xs2)));
2419            let closures: [_; N] = std::array::from_fn(|i| {
2420                let (c, a, o1, o2) = (&ctxs[i], &acts[i], p1[i], p2[i]);
2421                move |start: usize, end: usize| {
2422                    q8_range2_avx2(c.bytes, c.row_scale, &a.0, &a.1, c.cols, o1, o2, start, end)
2423                }
2424            });
2425            let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
2426                std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
2427            pool.run_many(&parts);
2428            return;
2429        }
2430        let closures: [_; N] = std::array::from_fn(|i| {
2431            let (c, o1, o2) = (&ctxs[i], p1[i], p2[i]);
2432            move |start: usize, end: usize| {
2433                q8_range2_f32(
2434                    c.bytes,
2435                    c.row_scale,
2436                    &c.xs1,
2437                    &c.xs2,
2438                    c.cols,
2439                    o1,
2440                    o2,
2441                    start,
2442                    end,
2443                )
2444            }
2445        });
2446        let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
2447            std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
2448        pool.run_many(&parts);
2449    }
2450
2451    /// Fused gate+up matvec with SiLU·mul: for each row r, computes
2452    /// `silu(gate·x) * (up·x)` and writes to `out[r]`. ONE pool dispatch,
2453    /// no intermediate g/u buffers, no separate silu pass. Falls back
2454    /// (returns false) for unsupported dtype combos.
2455    pub fn matvec_silu_mul(
2456        gate: &QTensor,
2457        up: &QTensor,
2458        x: &[f32],
2459        out: &mut [f32],
2460        pool: Option<&Pool>,
2461    ) -> bool {
2462        let inter = gate.rows();
2463        debug_assert_eq!(up.rows(), inter);
2464        debug_assert_eq!(out.len(), inter);
2465        debug_assert_eq!(gate.cols(), up.cols());
2466        if !a8w8_enabled() {
2467            return false;
2468        }
2469        let act = split_act(x);
2470        let act = &act;
2471        let x_ref = x;
2472        let out_addr = SendMut(out.as_mut_ptr());
2473
2474        match (gate, up) {
2475            // Q4Block gate + Q4Block up (most common mobile q4 models)
2476            (
2477                Self::Mapped {
2478                    dtype: TensorDtype::Q4Block,
2479                    ..
2480                },
2481                Self::Mapped {
2482                    dtype: TensorDtype::Q4Block,
2483                    ..
2484                },
2485            ) => {
2486                let (gp, gs) = q4_split(gate.quant_bytes(), gate.rows(), gate.cols());
2487                let (up_p, up_s) = q4_split(up.quant_bytes(), up.rows(), up.cols());
2488                let gpr = gate.cols() / GROUP_SIZE;
2489                let cols = gate.cols();
2490                let run = move |start: usize, end: usize| {
2491                    for r in start..end {
2492                        let mut gv = dot_q4_row_i8(gp, gs, r * gpr, gpr, &act.xq) * act.sx;
2493                        let mut uv = dot_q4_row_i8(up_p, up_s, r * gpr, gpr, &act.xq) * act.sx;
2494                        for &(j, xv) in &act.outliers {
2495                            let flat = r * cols + j;
2496                            let gb = gp[flat / 2];
2497                            let gn = if flat & 1 == 0 { gb & 0x0F } else { gb >> 4 };
2498                            let gsc = f16_to_f32(u16::from_le_bytes([
2499                                gs[(flat / GROUP_SIZE) * 2],
2500                                gs[(flat / GROUP_SIZE) * 2 + 1],
2501                            ]));
2502                            gv += ((gn as i32 - 8) as f32) * gsc * xv;
2503                            let ub = up_p[flat / 2];
2504                            let un = if flat & 1 == 0 { ub & 0x0F } else { ub >> 4 };
2505                            let usc = f16_to_f32(u16::from_le_bytes([
2506                                up_s[(flat / GROUP_SIZE) * 2],
2507                                up_s[(flat / GROUP_SIZE) * 2 + 1],
2508                            ]));
2509                            uv += ((un as i32 - 8) as f32) * usc * xv;
2510                        }
2511                        let silu_g = gv / (1.0 + (-gv).exp());
2512                        // SAFETY: disjoint row ranges per worker.
2513                        unsafe { *out_addr.at(r) = silu_g * uv };
2514                    }
2515                };
2516                dispatch_rows(pool, inter, &run);
2517                true
2518            }
2519            // Q4Tiled gate + Q4Tiled up — one row pass, both tile
2520            // streams sequential, silu·mul fused (same per-row math as
2521            // `q4t_matvec`).
2522            (
2523                Self::Mapped {
2524                    dtype: TensorDtype::Q4Tiled,
2525                    ..
2526                },
2527                Self::Mapped {
2528                    dtype: TensorDtype::Q4Tiled,
2529                    ..
2530                },
2531            ) => {
2532                let g_bytes = gate.quant_bytes();
2533                let u_bytes = up.quant_bytes();
2534                let gpr = gate.cols() / GROUP_SIZE;
2535                let run = move |start: usize, end: usize| {
2536                    for r in start..end {
2537                        let mut gv = dot_q4t_row_i8(g_bytes, r, gpr, &act.xq) * act.sx;
2538                        let mut uv = dot_q4t_row_i8(u_bytes, r, gpr, &act.xq) * act.sx;
2539                        for &(j, xv) in &act.outliers {
2540                            let (w, s) = q4t_outlier(g_bytes, r, gpr, j);
2541                            gv += w * s * xv;
2542                            let (w, s) = q4t_outlier(u_bytes, r, gpr, j);
2543                            uv += w * s * xv;
2544                        }
2545                        let silu_g = gv / (1.0 + (-gv).exp());
2546                        // SAFETY: disjoint row ranges per worker.
2547                        unsafe { *out_addr.at(r) = silu_g * uv };
2548                    }
2549                };
2550                dispatch_rows(pool, inter, &run);
2551                true
2552            }
2553            // Q4TiledP gate + Q4TiledP up — the same fused row pass, with
2554            // each row's two ladders built once and spent on both streams.
2555            (
2556                Self::Mapped {
2557                    dtype: TensorDtype::Q4TiledP,
2558                    ..
2559                },
2560                Self::Mapped {
2561                    dtype: TensorDtype::Q4TiledP,
2562                    ..
2563                },
2564            ) => {
2565                let cols = gate.cols();
2566                let gpr = cols / GROUP_SIZE;
2567                let gv_view = Q4tpView::new(gate.quant_bytes(), inter, cols);
2568                let uv_view = Q4tpView::new(up.quant_bytes(), inter, cols);
2569                let run = |start: usize, end: usize| {
2570                    let (mut gsc, mut usc) = (vec![0f32; gpr], vec![0f32; gpr]);
2571                    for r in start..end {
2572                        gv_view.scales_into(r, gpr, &mut gsc);
2573                        uv_view.scales_into(r, gpr, &mut usc);
2574                        let mut gv = dot_q4tp_row_i8(gv_view.nib, r, gpr, &act.xq, &gsc) * act.sx;
2575                        let mut uv = dot_q4tp_row_i8(uv_view.nib, r, gpr, &act.xq, &usc) * act.sx;
2576                        for &(j, xv) in &act.outliers {
2577                            let (w, s) = q4tp_outlier(gv_view.nib, r, gpr, j, &gsc);
2578                            gv += w * s * xv;
2579                            let (w, s) = q4tp_outlier(uv_view.nib, r, gpr, j, &usc);
2580                            uv += w * s * xv;
2581                        }
2582                        let silu_g = gv / (1.0 + (-gv).exp());
2583                        // SAFETY: disjoint row ranges per worker.
2584                        unsafe { *out_addr.at(r) = silu_g * uv };
2585                    }
2586                };
2587                dispatch_rows(pool, inter, &run);
2588                true
2589            }
2590            // Q1 gate + Q1 up — one row pass over both sign streams,
2591            // silu·mul fused (the per-row math of `q1_range_a8w8`); the
2592            // activation group sums are shared by both streams. Without
2593            // this arm a q1 dense FFN paid two dispatches + a combine
2594            // loop — the exact barrier this function exists to remove.
2595            (
2596                Self::Mapped {
2597                    dtype: TensorDtype::Q1,
2598                    ..
2599                },
2600                Self::Mapped {
2601                    dtype: TensorDtype::Q1,
2602                    ..
2603                },
2604            ) => {
2605                let g_bytes = gate.quant_bytes();
2606                let u_bytes = up.quant_bytes();
2607                let gpr = gate.cols() / GROUP_SIZE;
2608                let gsum = q1_group_sums(&act.xq, gpr);
2609                let gsum = &gsum;
2610                let run = move |start: usize, end: usize| {
2611                    for r in start..end {
2612                        let mut gv = dot_q1_row_i8(g_bytes, r, gpr, &act.xq, gsum) * act.sx;
2613                        let mut uv = dot_q1_row_i8(u_bytes, r, gpr, &act.xq, gsum) * act.sx;
2614                        for &(j, xv) in &act.outliers {
2615                            let (w, s) = q1_outlier(g_bytes, r, gpr, j);
2616                            gv += w * s * xv;
2617                            let (w, s) = q1_outlier(u_bytes, r, gpr, j);
2618                            uv += w * s * xv;
2619                        }
2620                        let silu_g = gv / (1.0 + (-gv).exp());
2621                        // SAFETY: disjoint row ranges per worker.
2622                        unsafe { *out_addr.at(r) = silu_g * uv };
2623                    }
2624                };
2625                dispatch_rows(pool, inter, &run);
2626                true
2627            }
2628            // Q2TiledP gate + Q2TiledP up — the 2-bit expert pair (MoE
2629            // FFNs of the W2 class): one row pass, both ladders built
2630            // once, integer code dots with shared group sums.
2631            (
2632                Self::Mapped {
2633                    dtype: TensorDtype::Q2TiledP,
2634                    ..
2635                },
2636                Self::Mapped {
2637                    dtype: TensorDtype::Q2TiledP,
2638                    ..
2639                },
2640            ) => {
2641                let cols = gate.cols();
2642                let gpr = cols / GROUP_SIZE;
2643                let gv_view = Q4tpView::new_q2(gate.quant_bytes(), inter, cols);
2644                let uv_view = Q4tpView::new_q2(up.quant_bytes(), inter, cols);
2645                let gsum = q1_group_sums(&act.xq, gpr);
2646                let gsum = &gsum;
2647                let run = move |start: usize, end: usize| {
2648                    let (mut gsc, mut usc) = (vec![0f32; gpr], vec![0f32; gpr]);
2649                    for r in start..end {
2650                        gv_view.scales_into(r, gpr, &mut gsc);
2651                        uv_view.scales_into(r, gpr, &mut usc);
2652                        let mut gv =
2653                            dot_q2tp_row_i8(gv_view.nib, r, gpr, &act.xq, gsum, &gsc) * act.sx;
2654                        let mut uv =
2655                            dot_q2tp_row_i8(uv_view.nib, r, gpr, &act.xq, gsum, &usc) * act.sx;
2656                        for &(j, xv) in &act.outliers {
2657                            let (w, s) = q2tp_outlier(gv_view.nib, r, gpr, j, &gsc);
2658                            gv += w * s * xv;
2659                            let (w, s) = q2tp_outlier(uv_view.nib, r, gpr, j, &usc);
2660                            uv += w * s * xv;
2661                        }
2662                        let silu_g = gv / (1.0 + (-gv).exp());
2663                        // SAFETY: disjoint row ranges per worker.
2664                        unsafe { *out_addr.at(r) = silu_g * uv };
2665                    }
2666                };
2667                dispatch_rows(pool, inter, &run);
2668                true
2669            }
2670            // Q8Row gate + Q8Row up — one row pass over both i8 streams.
2671            // Q8_2f stays out on purpose: its column field prescales the
2672            // activations PER TENSOR, which breaks this fn's shared
2673            // split_act contract — it keeps the two-dispatch path.
2674            (
2675                Self::Mapped {
2676                    dtype: TensorDtype::Q8Row,
2677                    row_scale: g_rs,
2678                    ..
2679                },
2680                Self::Mapped {
2681                    dtype: TensorDtype::Q8Row,
2682                    row_scale: u_rs,
2683                    ..
2684                },
2685            ) => {
2686                let g_bytes = gate.quant_bytes();
2687                let u_bytes = up.quant_bytes();
2688                let cols = gate.cols();
2689                let run = move |start: usize, end: usize| {
2690                    for r in start..end {
2691                        let gv = q8_row_dot(&g_bytes[r * cols..(r + 1) * cols], act) * g_rs[r];
2692                        let uv = q8_row_dot(&u_bytes[r * cols..(r + 1) * cols], act) * u_rs[r];
2693                        let silu_g = gv / (1.0 + (-gv).exp());
2694                        // SAFETY: disjoint row ranges per worker.
2695                        unsafe { *out_addr.at(r) = silu_g * uv };
2696                    }
2697                };
2698                dispatch_rows(pool, inter, &run);
2699                true
2700            }
2701            // Q1T gate + Q1T up
2702            (
2703                Self::Mapped {
2704                    dtype: TensorDtype::Q1T,
2705                    ..
2706                },
2707                Self::Mapped {
2708                    dtype: TensorDtype::Q1T,
2709                    ..
2710                },
2711            ) => {
2712                const TILE: usize = cortiq_core::quant::Q1T_TILE;
2713                let g_bytes = gate.quant_bytes();
2714                let u_bytes = up.quant_bytes();
2715                let gpr = gate.cols() / GROUP_SIZE;
2716                let (g_rp, g_ent, g_ov) = q1t_overlay(g_bytes, inter * gpr * TILE, inter);
2717                let (u_rp, u_ent, u_ov) = q1t_overlay(u_bytes, inter * gpr * TILE, inter);
2718                let run = move |start: usize, end: usize| {
2719                    for r in start..end {
2720                        let mut gv = q1t_dot_row_i8(g_bytes, r, gpr, &act.xq) * act.sx;
2721                        let mut uv = q1t_dot_row_i8(u_bytes, r, gpr, &act.xq) * act.sx;
2722                        for &(j, xv) in &act.outliers {
2723                            gv += q1t_base_weight(g_bytes, r, gpr, j) * xv;
2724                            uv += q1t_base_weight(u_bytes, r, gpr, j) * xv;
2725                        }
2726                        gv += q1t_row_outlier_correction(g_bytes, r, g_rp, g_ent, g_ov, x_ref);
2727                        uv += q1t_row_outlier_correction(u_bytes, r, u_rp, u_ent, u_ov, x_ref);
2728                        let silu_g = gv / (1.0 + (-gv).exp());
2729                        // SAFETY: disjoint row ranges per worker.
2730                        unsafe { *out_addr.at(r) = silu_g * uv };
2731                    }
2732                };
2733                dispatch_rows(pool, inter, &run);
2734                true
2735            }
2736            _ => false,
2737        }
2738    }
2739
2740    /// Every routed expert's fused gate/up/SiLU under ONE pool dispatch.
2741    ///
2742    /// The per-expert path pays a pool barrier per expert per stage: at 9
2743    /// experts over 40 layers that is ~720 barriers a token, and a decode
2744    /// profile of Qwen3.6-35B-A3B showed the pool parked in
2745    /// `psynch_cvwait` about twice as long as it spent computing. Laying
2746    /// every expert's rows end-to-end in one virtual row space collapses
2747    /// the stage to a single dispatch. The per-row body is the
2748    /// single-expert q4tp arm verbatim, so outputs are bit-identical.
2749    ///
2750    /// `false` = something is outside the fused q4tp kernel (dtype, shape,
2751    /// or the `CMF_SDOT=0` exact contract); the caller walks the ordinary
2752    /// per-expert path.
2753    pub fn moe_gate_up_many(
2754        pairs: &[(&QTensor, &QTensor)],
2755        x: &[f32],
2756        outs: &mut [Vec<f32>],
2757        pool: Option<&Pool>,
2758    ) -> bool {
2759        if pairs.is_empty() || pairs.len() != outs.len() || !a8w8_enabled() {
2760            return false;
2761        }
2762        let inter = pairs[0].0.rows();
2763        let cols = pairs[0].0.cols();
2764        if cols % GROUP_SIZE != 0 {
2765            return false;
2766        }
2767        let gpr = cols / GROUP_SIZE;
2768        // Uniform layout across every routed pair: q4tp, or the 2-bit
2769        // profile's q2tp gate/up (the W2 class). Mixed sets refuse.
2770        let q2 = matches!(
2771            pairs[0].0,
2772            Self::Mapped {
2773                dtype: TensorDtype::Q2TiledP,
2774                ..
2775            }
2776        );
2777        let want = if q2 {
2778            TensorDtype::Q2TiledP
2779        } else {
2780            TensorDtype::Q4TiledP
2781        };
2782        let mut views = Vec::with_capacity(pairs.len() * 2);
2783        for ((g, u), o) in pairs.iter().zip(outs.iter()) {
2784            let both = matches!(g, Self::Mapped { dtype, .. } if *dtype == want)
2785                && matches!(u, Self::Mapped { dtype, .. } if *dtype == want);
2786            if !both
2787                || g.rows() != inter
2788                || u.rows() != inter
2789                || g.cols() != cols
2790                || u.cols() != cols
2791                || o.len() != inter
2792            {
2793                return false;
2794            }
2795            let mk = if q2 { Q4tpView::new_q2 } else { Q4tpView::new };
2796            views.push(mk(g.quant_bytes(), inter, cols));
2797            views.push(mk(u.quant_bytes(), inter, cols));
2798        }
2799        let act = split_act(x);
2800        let gsum = if q2 {
2801            q1_group_sums(&act.xq, gpr)
2802        } else {
2803            Vec::new()
2804        };
2805        let (act, gsum) = (&act, &gsum);
2806        let ptrs: Vec<SendMut> = outs.iter_mut().map(|o| SendMut(o.as_mut_ptr())).collect();
2807        let (views, ptrs) = (&views, &ptrs);
2808        let run = |start: usize, end: usize| {
2809            let (mut gsc, mut usc) = (vec![0f32; gpr], vec![0f32; gpr]);
2810            for flat in start..end {
2811                let (e, r) = (flat / inter, flat % inter);
2812                let gv_view = &views[e * 2];
2813                let uv_view = &views[e * 2 + 1];
2814                gv_view.scales_into(r, gpr, &mut gsc);
2815                uv_view.scales_into(r, gpr, &mut usc);
2816                let (mut gv, mut uv) = if q2 {
2817                    (
2818                        dot_q2tp_row_i8(gv_view.nib, r, gpr, &act.xq, gsum, &gsc) * act.sx,
2819                        dot_q2tp_row_i8(uv_view.nib, r, gpr, &act.xq, gsum, &usc) * act.sx,
2820                    )
2821                } else {
2822                    (
2823                        dot_q4tp_row_i8(gv_view.nib, r, gpr, &act.xq, &gsc) * act.sx,
2824                        dot_q4tp_row_i8(uv_view.nib, r, gpr, &act.xq, &usc) * act.sx,
2825                    )
2826                };
2827                for &(j, xv) in &act.outliers {
2828                    let (og, ou) = if q2 {
2829                        (
2830                            q2tp_outlier(gv_view.nib, r, gpr, j, &gsc),
2831                            q2tp_outlier(uv_view.nib, r, gpr, j, &usc),
2832                        )
2833                    } else {
2834                        (
2835                            q4tp_outlier(gv_view.nib, r, gpr, j, &gsc),
2836                            q4tp_outlier(uv_view.nib, r, gpr, j, &usc),
2837                        )
2838                    };
2839                    gv += og.0 * og.1 * xv;
2840                    uv += ou.0 * ou.1 * xv;
2841                }
2842                let silu_g = gv / (1.0 + (-gv).exp());
2843                // SAFETY: one worker owns each (expert, row) pair.
2844                unsafe { *ptrs[e].at(r) = silu_g * uv };
2845            }
2846        };
2847        dispatch_rows(pool, pairs.len() * inter, &run);
2848        true
2849    }
2850
2851    /// Every routed expert's down projection, weighted and summed into
2852    /// `out`, under ONE pool dispatch.
2853    ///
2854    /// Partitioned by OUTPUT row rather than by expert: each row is owned
2855    /// by a single worker, so the experts are summed in the caller's order
2856    /// — the same sequence of f32 adds the serial `out[i] += w·eo[i]` loop
2857    /// performs, hence bit-identical. Partitioning by expert instead would
2858    /// race on the shared accumulator.
2859    pub fn moe_down_many(
2860        downs: &[&QTensor],
2861        gs: &[Vec<f32>],
2862        weights: &[f32],
2863        out: &mut [f32],
2864        pool: Option<&Pool>,
2865    ) -> bool {
2866        if downs.is_empty()
2867            || downs.len() != gs.len()
2868            || downs.len() != weights.len()
2869            || !a8w8_enabled()
2870        {
2871            return false;
2872        }
2873        let rows = out.len();
2874        let cols = downs[0].cols();
2875        if cols % GROUP_SIZE != 0 {
2876            return false;
2877        }
2878        let gpr = cols / GROUP_SIZE;
2879        let mut views = Vec::with_capacity(downs.len());
2880        for (d, g) in downs.iter().zip(gs.iter()) {
2881            if !matches!(
2882                d,
2883                Self::Mapped {
2884                    dtype: TensorDtype::Q4TiledP,
2885                    ..
2886                }
2887            ) || d.rows() != rows
2888                || d.cols() != cols
2889                || g.len() != cols
2890            {
2891                return false;
2892            }
2893            views.push(Q4tpView::new(d.quant_bytes(), rows, cols));
2894        }
2895        // One int8 split per expert — the activation vectors differ.
2896        let acts: Vec<SplitAct> = gs.iter().map(|g| split_act(g)).collect();
2897        // Partitioned by OUTPUT row, with the experts folded inside: each
2898        // row is owned by one worker, so they are summed in the caller's
2899        // order — the same f32 sequence the serial `out[i] += w·eo[i]`
2900        // loop produces. Partitioning by expert instead would either race
2901        // on the accumulator or need a scratch plane and a second pass;
2902        // measured, that variant was a wash, so this keeps the simpler
2903        // shape.
2904        let out_addr = SendMut(out.as_mut_ptr());
2905        let (views, acts, weights) = (&views, &acts, &weights);
2906        let run = |start: usize, end: usize| {
2907            let mut sc = vec![0f32; gpr];
2908            for r in start..end {
2909                let mut acc = 0f32;
2910                for (e, v) in views.iter().enumerate() {
2911                    v.scales_into(r, gpr, &mut sc);
2912                    let a = &acts[e];
2913                    let mut d = dot_q4tp_row_i8(v.nib, r, gpr, &a.xq, &sc) * a.sx;
2914                    for &(j, xv) in &a.outliers {
2915                        let (w, s) = q4tp_outlier(v.nib, r, gpr, j, &sc);
2916                        d += w * s * xv;
2917                    }
2918                    acc += weights[e] * d;
2919                }
2920                // SAFETY: disjoint row ranges per worker.
2921                unsafe { *out_addr.at(r) = acc };
2922            }
2923        };
2924        dispatch_rows(pool, rows, &run);
2925        true
2926    }
2927}
2928
2929/// Batched q8 kernel: same math as qmatvec, the row makes a single
2930/// pass from memory for the whole batch.
2931/// Accelerate CBLAS — the Apple AMX matrix units, the same engine
2932/// llama.cpp's `-ngl 0` prefill rides via ggml-blas.
2933#[cfg(target_os = "macos")]
2934mod accel_blas {
2935    #[link(name = "Accelerate", kind = "framework")]
2936    unsafe extern "C" {
2937        pub fn cblas_sgemm(
2938            order: i32,
2939            trans_a: i32,
2940            trans_b: i32,
2941            m: i32,
2942            n: i32,
2943            k: i32,
2944            alpha: f32,
2945            a: *const f32,
2946            lda: i32,
2947            b: *const f32,
2948            ldb: i32,
2949            beta: f32,
2950            c: *mut f32,
2951            ldc: i32,
2952        );
2953    }
2954}
2955
2956#[cfg(target_os = "macos")]
2957pub(crate) fn accel_gemm_enabled() -> bool {
2958    static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
2959    *ON.get_or_init(|| std::env::var("CMF_ACCEL").map(|v| v != "0").unwrap_or(true))
2960}
2961
2962/// Off macOS the "accel" GEMM is the portable NEON micro-kernel below —
2963/// same entry point, so the batched-attention path opens on mobile.
2964#[cfg(all(target_arch = "aarch64", not(target_os = "macos")))]
2965pub(crate) fn accel_gemm_enabled() -> bool {
2966    true
2967}
2968
2969/// Portable NEON f32 GEMM (row-major, optional Bᵀ): a 4×8 fmla
2970/// micro-kernel with A broadcast against B panels — the mobile stand-in
2971/// for Accelerate in the batched causal attention (QKᵀ and P·V). Not a
2972/// BLAS: shapes here are the attention panels (m ≤ heads·chunk,
2973/// k = head_dim or context), and the goal is removing the per-position
2974/// quadratic wall, not peak GEMM.
2975#[cfg(target_arch = "aarch64")]
2976#[allow(clippy::too_many_arguments)]
2977pub(crate) fn neon_gemm_rm(
2978    m: usize,
2979    n: usize,
2980    k: usize,
2981    alpha: f32,
2982    a: &[f32],
2983    lda: usize,
2984    b_mat: &[f32],
2985    ldb: usize,
2986    b_rows_are_n: bool,
2987    c: &mut [f32],
2988    ldc: usize,
2989) {
2990    debug_assert!(a.len() >= (m - 1) * lda + k);
2991    debug_assert!(c.len() >= (m - 1) * ldc + n);
2992    // SAFETY: bounds asserted above; NEON is baseline on aarch64.
2993    unsafe {
2994        use core::arch::aarch64::*;
2995        let mut i = 0usize;
2996        while i < m {
2997            let mi = (m - i).min(4);
2998            let mut j = 0usize;
2999            while j < n {
3000                let nj = (n - j).min(8);
3001                if mi == 4 && nj == 8 {
3002                    let (mut c0a, mut c0b) = (vdupq_n_f32(0.0), vdupq_n_f32(0.0));
3003                    let (mut c1a, mut c1b) = (vdupq_n_f32(0.0), vdupq_n_f32(0.0));
3004                    let (mut c2a, mut c2b) = (vdupq_n_f32(0.0), vdupq_n_f32(0.0));
3005                    let (mut c3a, mut c3b) = (vdupq_n_f32(0.0), vdupq_n_f32(0.0));
3006                    for p in 0..k {
3007                        let (b0, b1) = if b_rows_are_n {
3008                            // B is [n, k]: column p of Bᵀ = element p of
3009                            // eight consecutive B rows — gathered.
3010                            let base = b_mat.as_ptr().add(j * ldb + p);
3011                            let g = |o: usize| *base.add(o * ldb);
3012                            ([g(0), g(1), g(2), g(3)], [g(4), g(5), g(6), g(7)])
3013                        } else {
3014                            let base = b_mat.as_ptr().add(p * ldb + j);
3015                            (
3016                                [*base, *base.add(1), *base.add(2), *base.add(3)],
3017                                [*base.add(4), *base.add(5), *base.add(6), *base.add(7)],
3018                            )
3019                        };
3020                        let bv0 = vld1q_f32(b0.as_ptr());
3021                        let bv1 = vld1q_f32(b1.as_ptr());
3022                        let a0 = vdupq_n_f32(*a.as_ptr().add(i * lda + p));
3023                        let a1 = vdupq_n_f32(*a.as_ptr().add((i + 1) * lda + p));
3024                        let a2 = vdupq_n_f32(*a.as_ptr().add((i + 2) * lda + p));
3025                        let a3 = vdupq_n_f32(*a.as_ptr().add((i + 3) * lda + p));
3026                        c0a = vfmaq_f32(c0a, a0, bv0);
3027                        c0b = vfmaq_f32(c0b, a0, bv1);
3028                        c1a = vfmaq_f32(c1a, a1, bv0);
3029                        c1b = vfmaq_f32(c1b, a1, bv1);
3030                        c2a = vfmaq_f32(c2a, a2, bv0);
3031                        c2b = vfmaq_f32(c2b, a2, bv1);
3032                        c3a = vfmaq_f32(c3a, a3, bv0);
3033                        c3b = vfmaq_f32(c3b, a3, bv1);
3034                    }
3035                    let al = vdupq_n_f32(alpha);
3036                    for (r, (ca, cb)) in [(c0a, c0b), (c1a, c1b), (c2a, c2b), (c3a, c3b)]
3037                        .iter()
3038                        .enumerate()
3039                    {
3040                        let dst = c.as_mut_ptr().add((i + r) * ldc + j);
3041                        vst1q_f32(dst, vmulq_f32(*ca, al));
3042                        vst1q_f32(dst.add(4), vmulq_f32(*cb, al));
3043                    }
3044                } else {
3045                    for r in 0..mi {
3046                        for q in 0..nj {
3047                            let mut acc = 0f32;
3048                            for p in 0..k {
3049                                let bv = if b_rows_are_n {
3050                                    b_mat[(j + q) * ldb + p]
3051                                } else {
3052                                    b_mat[p * ldb + j + q]
3053                                };
3054                                acc += a[(i + r) * lda + p] * bv;
3055                            }
3056                            c[(i + r) * ldc + j + q] = acc * alpha;
3057                        }
3058                    }
3059                }
3060                j += nj;
3061            }
3062            i += mi;
3063        }
3064    }
3065}
3066
3067/// Off-macOS aarch64: the batched attention rides the NEON micro-GEMM.
3068#[cfg(all(target_arch = "aarch64", not(target_os = "macos")))]
3069#[allow(clippy::too_many_arguments)]
3070pub(crate) fn sgemm_rm(
3071    m: usize,
3072    n: usize,
3073    k: usize,
3074    alpha: f32,
3075    a: &[f32],
3076    lda: usize,
3077    b_mat: &[f32],
3078    ldb: usize,
3079    b_rows_are_n: bool,
3080    c: &mut [f32],
3081    ldc: usize,
3082) {
3083    neon_gemm_rm(m, n, k, alpha, a, lda, b_mat, ldb, b_rows_are_n, c, ldc);
3084}
3085
3086/// Row-major f32 GEMM, exposed for offline tools (the AWNP pass builds a
3087/// per-layer projection and applies it to every expert; a naive triple loop
3088/// would turn a two-minute job into half an hour).
3089#[allow(clippy::too_many_arguments)]
3090pub fn sgemm_public(
3091    m: usize,
3092    n: usize,
3093    k: usize,
3094    alpha: f32,
3095    a: &[f32],
3096    lda: usize,
3097    b_mat: &[f32],
3098    ldb: usize,
3099    b_rows_are_n: bool,
3100    c: &mut [f32],
3101    ldc: usize,
3102) {
3103    #[cfg(any(target_os = "macos", target_arch = "aarch64"))]
3104    {
3105        sgemm_rm(m, n, k, alpha, a, lda, b_mat, ldb, b_rows_are_n, c, ldc);
3106    }
3107    // x86 without Accelerate has no sgemm_rm: the specialized paths there are
3108    // quantized kernels, not an f32 GEMM. Only the offline AWNP pass reaches
3109    // this, so correctness matters and throughput does not — a triple loop is
3110    // the honest fallback rather than a reason to make the tool macOS-only.
3111    #[cfg(not(any(target_os = "macos", target_arch = "aarch64")))]
3112    {
3113        for i in 0..m {
3114            for j in 0..n {
3115                let mut acc = 0f32;
3116                for p in 0..k {
3117                    let bv = if b_rows_are_n {
3118                        b_mat[j * ldb + p]
3119                    } else {
3120                        b_mat[p * ldb + j]
3121                    };
3122                    acc += a[i * lda + p] * bv;
3123                }
3124                c[i * ldc + j] = alpha * acc;
3125            }
3126        }
3127    }
3128}
3129
3130/// Row-major f32 GEMM on Accelerate: C[m,n] = alpha·A[m,k] × B(ᵀ).
3131/// `b_rows_are_n` = true multiplies by Bᵀ where B is stored [n, k].
3132#[cfg(target_os = "macos")]
3133#[allow(clippy::too_many_arguments)]
3134pub(crate) fn sgemm_rm(
3135    m: usize,
3136    n: usize,
3137    k: usize,
3138    alpha: f32,
3139    a: &[f32],
3140    lda: usize,
3141    b_mat: &[f32],
3142    ldb: usize,
3143    b_rows_are_n: bool,
3144    c: &mut [f32],
3145    ldc: usize,
3146) {
3147    debug_assert!(a.len() >= (m - 1) * lda + k);
3148    debug_assert!(c.len() >= (m - 1) * ldc + n);
3149    // Test hook: route the attention GEMMs through the portable NEON
3150    // micro-kernel ON APPLE SILICON — how the mobile batched attend is
3151    // measured without a phone in the loop. (Intel macOS has no NEON —
3152    // the hook is a no-op there, Accelerate continues below.)
3153    #[cfg(target_arch = "aarch64")]
3154    if std::env::var("CMF_FORCE_NEON_GEMM")
3155        .map(|v| v == "1")
3156        .unwrap_or(false)
3157    {
3158        return neon_gemm_rm(m, n, k, alpha, a, lda, b_mat, ldb, b_rows_are_n, c, ldc);
3159    }
3160    unsafe {
3161        accel_blas::cblas_sgemm(
3162            101, // RowMajor
3163            111, // NoTrans A
3164            if b_rows_are_n { 112 } else { 111 },
3165            m as i32,
3166            n as i32,
3167            k as i32,
3168            alpha,
3169            a.as_ptr(),
3170            lda as i32,
3171            b_mat.as_ptr(),
3172            ldb as i32,
3173            0.0,
3174            c.as_mut_ptr(),
3175            ldc as i32,
3176        );
3177    }
3178}
3179
3180/// Prefill GEMM through Accelerate (macOS): dequantize q8 rows into
3181/// f32 tiles (scale folded in, pool-parallel) and multiply each tile
3182/// on the AMX with one row-major sgemm. Tiles live in cache, weights
3183/// stream once. Numerics are f32-GEMM (not the int8 dot): prefill
3184/// logits shift within f32 rounding — tolerance-class, like every
3185/// reduction-order change; decode (M=1) never takes this path.
3186#[cfg(target_os = "macos")]
3187fn qmatmat_accel(
3188    q: &[u8],
3189    row_scale: &[f32],
3190    pre: &[std::borrow::Cow<'_, [f32]>],
3191    rows: usize,
3192    cols: usize,
3193    out: &mut [f32],
3194    pool: Option<&Pool>,
3195) {
3196    // NOTE: double-buffering the dequant against the sgemm (a scoped
3197    // thread driving the pool on tile k+1 while the caller multiplies
3198    // tile k) was tried and LOST ~6%: Accelerate's sgemm is itself
3199    // multithreaded, and the dequant workers just steal its cores.
3200    const TR: usize = 2048;
3201    let b = pre.len();
3202    thread_local! {
3203        static XPANEL: std::cell::RefCell<Vec<f32>> = const { std::cell::RefCell::new(Vec::new()) };
3204        static WTILE: std::cell::RefCell<Vec<f32>> = const { std::cell::RefCell::new(Vec::new()) };
3205    }
3206    XPANEL.with(|xp| {
3207        WTILE.with(|wt| {
3208            let mut xpanel = xp.borrow_mut();
3209            xpanel.clear();
3210            for x in pre {
3211                xpanel.extend_from_slice(x);
3212            }
3213            let mut wtile = wt.borrow_mut();
3214            wtile.resize(TR * cols, 0.0);
3215            let mut r0 = 0usize;
3216            while r0 < rows {
3217                let tr = TR.min(rows - r0);
3218                // Dequant the tile (scale folded) — pool-parallel.
3219                let wt_addr = SendMut(wtile.as_mut_ptr());
3220                let run = |start: usize, end: usize| {
3221                    for r in start..end {
3222                        let row = &q[(r0 + r) * cols..(r0 + r + 1) * cols];
3223                        let s = row_scale[r0 + r];
3224                        // SAFETY: workers cover disjoint r ranges.
3225                        let dst =
3226                            unsafe { std::slice::from_raw_parts_mut(wt_addr.at(r * cols), cols) };
3227                        for (d, &v) in dst.iter_mut().zip(row) {
3228                            *d = (v as i8) as f32 * s;
3229                        }
3230                    }
3231                };
3232                dispatch_rows(pool, tr, &run);
3233                // C[b, tr] (at column r0 of out[b, rows]) = X · Wtileᵀ
3234                unsafe {
3235                    accel_blas::cblas_sgemm(
3236                        101, // RowMajor
3237                        111, // NoTrans A
3238                        112, // Trans B
3239                        b as i32,
3240                        tr as i32,
3241                        cols as i32,
3242                        1.0,
3243                        xpanel.as_ptr(),
3244                        cols as i32,
3245                        wtile.as_ptr(),
3246                        cols as i32,
3247                        0.0,
3248                        out.as_mut_ptr().add(r0),
3249                        rows as i32,
3250                    );
3251                }
3252                r0 += tr;
3253            }
3254        })
3255    });
3256}
3257
3258fn qmatmat(
3259    q: &[u8],
3260    row_scale: &[f32],
3261    pre: &[std::borrow::Cow<'_, [f32]>],
3262    rows: usize,
3263    cols: usize,
3264    out: &mut [f32],
3265    pool: Option<&Pool>,
3266) {
3267    let b = pre.len();
3268    debug_assert_eq!(out.len(), b * rows);
3269    // Big prefill batches ride the AMX (roadmap PR3): the row×batch
3270    // SDOT loop below peaks near the CPU's dot throughput, an order
3271    // below the matrix units. Small tensors and tiny test models stay
3272    // on the exact integer path.
3273    #[cfg(target_os = "macos")]
3274    if b >= 8 && rows * cols >= 500_000 && accel_gemm_enabled() {
3275        qmatmat_accel(q, row_scale, pre, rows, cols, out, pool);
3276        return;
3277    }
3278    #[cfg(target_arch = "aarch64")]
3279    if sdot_enabled() {
3280        let acts: Vec<SplitAct> = pre.iter().map(|x| split_act(x)).collect();
3281        let out_addr = SendMut(out.as_mut_ptr());
3282        // Blocked 2×4 (mobile prefill: no AMX to fall back on — this
3283        // path IS the ARM prefill GEMM off Apple silicon).
3284        let blocked_ok = blocked_enabled();
3285        let use_i8mm = i8mm_enabled();
3286        if blocked_ok {
3287            let run = |start: usize, end: usize| {
3288                let mut o = start;
3289                while o < end {
3290                    if o + 2 <= end {
3291                        let r0 = &q[o * cols..(o + 1) * cols];
3292                        let r1 = &q[(o + 1) * cols..(o + 2) * cols];
3293                        let mut bi = 0usize;
3294                        while bi + 4 <= acts.len() {
3295                            let xs = [
3296                                acts[bi].xq.as_slice(),
3297                                acts[bi + 1].xq.as_slice(),
3298                                acts[bi + 2].xq.as_slice(),
3299                                acts[bi + 3].xq.as_slice(),
3300                            ];
3301                            let d = if use_i8mm {
3302                                unsafe { dot_i8_smmla_2x4(r0, r1, xs) }
3303                            } else {
3304                                unsafe { dot_i8_sdot_2x4(r0, r1, xs) }
3305                            };
3306                            for (r, row) in [r0, r1].into_iter().enumerate() {
3307                                for k in 0..4 {
3308                                    let act = &acts[bi + k];
3309                                    let mut v = d[r][k] as f32 * act.sx;
3310                                    for &(j, xv) in &act.outliers {
3311                                        v += (row[j] as i8) as f32 * xv;
3312                                    }
3313                                    unsafe {
3314                                        *out_addr.at((bi + k) * rows + o + r) = v * row_scale[o + r]
3315                                    };
3316                                }
3317                            }
3318                            bi += 4;
3319                        }
3320                        while bi < acts.len() {
3321                            for (r, row) in [r0, r1].into_iter().enumerate() {
3322                                let v = row_dot_sdot(row, &acts[bi]) * row_scale[o + r];
3323                                unsafe { *out_addr.at(bi * rows + o + r) = v };
3324                            }
3325                            bi += 1;
3326                        }
3327                        o += 2;
3328                    } else {
3329                        let row = &q[o * cols..(o + 1) * cols];
3330                        for (bi, act) in acts.iter().enumerate() {
3331                            let v = row_dot_sdot(row, act) * row_scale[o];
3332                            unsafe { *out_addr.at(bi * rows + o) = v };
3333                        }
3334                        o += 1;
3335                    }
3336                }
3337            };
3338            dispatch_rows(pool, rows, &run);
3339            return;
3340        }
3341        let run = |start: usize, end: usize| {
3342            for o in start..end {
3343                let row = &q[o * cols..(o + 1) * cols];
3344                for (bi, act) in acts.iter().enumerate() {
3345                    let v = row_dot_sdot(row, act) * row_scale[o];
3346                    unsafe { *out_addr.at(bi * rows + o) = v };
3347                }
3348            }
3349        };
3350        dispatch_rows(pool, rows, &run);
3351        return;
3352    }
3353    // x86 A8W8 batch. Non-VNNI parts take the BLOCKED 2×4 kernel
3354    // (roadmap P0: two weight rows' abs() stay in registers across four
3355    // activation streams); VNNI machines keep the per-row bias-trick
3356    // dot, which is already throughput-bound there.
3357    #[cfg(target_arch = "x86_64")]
3358    if avx2_a8w8_enabled() {
3359        let acts: Vec<SplitAct> = pre.iter().map(|x| split_act(x)).collect();
3360        let out_addr = SendMut(out.as_mut_ptr());
3361        // CMF_X86_BLOCKED=0 forces the per-row path (paired in-process
3362        // A/B on noisy shared-vCPU hosts).
3363        let blocked_ok = blocked_enabled();
3364        if !avx512vnni_enabled() && blocked_ok {
3365            let run = |start: usize, end: usize| {
3366                let mut o = start;
3367                while o < end {
3368                    if o + 2 <= end {
3369                        let r0 = &q[o * cols..(o + 1) * cols];
3370                        let r1 = &q[(o + 1) * cols..(o + 2) * cols];
3371                        let mut bi = 0usize;
3372                        while bi + 4 <= acts.len() {
3373                            let xs = [
3374                                acts[bi].xq.as_slice(),
3375                                acts[bi + 1].xq.as_slice(),
3376                                acts[bi + 2].xq.as_slice(),
3377                                acts[bi + 3].xq.as_slice(),
3378                            ];
3379                            let d = unsafe { dot_i8_i8_avx2_2x4(r0, r1, xs) };
3380                            for (r, row) in [r0, r1].into_iter().enumerate() {
3381                                for k in 0..4 {
3382                                    let act = &acts[bi + k];
3383                                    let mut v = d[r][k] as f32 * act.sx;
3384                                    for &(j, xv) in &act.outliers {
3385                                        v += (row[j] as i8) as f32 * xv;
3386                                    }
3387                                    unsafe {
3388                                        *out_addr.at((bi + k) * rows + o + r) = v * row_scale[o + r]
3389                                    };
3390                                }
3391                            }
3392                            bi += 4;
3393                        }
3394                        while bi < acts.len() {
3395                            for (r, row) in [r0, r1].into_iter().enumerate() {
3396                                let v = row_dot_avx2(row, &acts[bi]) * row_scale[o + r];
3397                                unsafe { *out_addr.at(bi * rows + o + r) = v };
3398                            }
3399                            bi += 1;
3400                        }
3401                        o += 2;
3402                    } else {
3403                        let row = &q[o * cols..(o + 1) * cols];
3404                        for (bi, act) in acts.iter().enumerate() {
3405                            let v = row_dot_avx2(row, act) * row_scale[o];
3406                            unsafe { *out_addr.at(bi * rows + o) = v };
3407                        }
3408                        o += 1;
3409                    }
3410                }
3411            };
3412            dispatch_rows(pool, rows, &run);
3413            return;
3414        }
3415        let run = |start: usize, end: usize| {
3416            for o in start..end {
3417                let row = &q[o * cols..(o + 1) * cols];
3418                for (bi, act) in acts.iter().enumerate() {
3419                    let v = row_dot_avx2(row, act) * row_scale[o];
3420                    unsafe { *out_addr.at(bi * rows + o) = v };
3421                }
3422            }
3423        };
3424        dispatch_rows(pool, rows, &run);
3425        return;
3426    }
3427    let out_addr = SendMut(out.as_mut_ptr());
3428    let run = |start: usize, end: usize| {
3429        for o in start..end {
3430            let row = &q[o * cols..(o + 1) * cols];
3431            for (bi, x) in pre.iter().enumerate() {
3432                let mut acc = 0f32;
3433                for j in 0..cols {
3434                    acc += (row[j] as i8) as f32 * x[j];
3435                }
3436                unsafe { *out_addr.at(bi * rows + o) = acc * row_scale[o] };
3437            }
3438        }
3439    };
3440    dispatch_rows(pool, rows, &run);
3441}
3442
3443/// Split rows across pool workers (shared qmatvec pattern). Self-balancing
3444/// — see `Pool::run_rows` for why a static 1/n split is wrong here.
3445fn dispatch_rows(pool: Option<&Pool>, rows: usize, run: &(dyn Fn(usize, usize) + Sync)) {
3446    match pool {
3447        Some(pool) if rows >= 256 => pool.run_rows(rows, run),
3448        _ => run(0, rows),
3449    }
3450}
3451
3452/// Split a q4_block blob into (packed nibbles, f16 group scales).
3453fn q4_split(bytes: &[u8], rows: usize, cols: usize) -> (&[u8], &[u8]) {
3454    let groups = rows * cols / GROUP_SIZE;
3455    bytes.split_at(groups * 16)
3456}
3457
3458/// SIMD unpack for the dominant vbit width B=4 (94% of rows on the
3459/// log2-shape calibration): 16 packed bytes -> 32 centered i8 values.
3460/// vbit packs MSB-first, so the HIGH nibble is the even element
3461/// (opposite of q4_block's lo-first interleave). Centering is u-7.
3462#[inline]
3463fn vbit_fill4(data: &[u8], buf: &mut [u8]) {
3464    #[cfg(target_arch = "aarch64")]
3465    unsafe {
3466        return vbit_fill4_neon(data, buf);
3467    }
3468    #[cfg(target_arch = "x86_64")]
3469    if avx2_enabled() {
3470        return unsafe { vbit_fill4_avx2(data, buf) };
3471    }
3472    #[allow(unreachable_code)]
3473    for (blk, chunk) in buf.chunks_exact_mut(8).enumerate() {
3474        let u = unpack8::<4>(&data[blk * 4..]);
3475        for k in 0..8 {
3476            chunk[k] = (u[k] - 7) as i8 as u8;
3477        }
3478    }
3479}
3480
3481#[cfg(target_arch = "aarch64")]
3482#[target_feature(enable = "neon")]
3483unsafe fn vbit_fill4_neon(data: &[u8], buf: &mut [u8]) {
3484    // SAFETY: buf.len() is a multiple of GROUP_SIZE=32; data holds
3485    // buf.len()/2 packed bytes (validated at load).
3486    unsafe {
3487        use core::arch::aarch64::*;
3488        let n = buf.len();
3489        let mask = vdupq_n_u8(0x0F);
3490        let seven = vdupq_n_s8(7);
3491        let mut g = 0usize;
3492        while g * 32 + 32 <= n {
3493            let b = vld1q_u8(data.as_ptr().add(g * 16));
3494            let hi = vshrq_n_u8::<4>(b);
3495            let lo = vandq_u8(b, mask);
3496            let z0 = vsubq_s8(vreinterpretq_s8_u8(vzip1q_u8(hi, lo)), seven);
3497            let z1 = vsubq_s8(vreinterpretq_s8_u8(vzip2q_u8(hi, lo)), seven);
3498            vst1q_u8(buf.as_mut_ptr().add(g * 32), vreinterpretq_u8_s8(z0));
3499            vst1q_u8(buf.as_mut_ptr().add(g * 32 + 16), vreinterpretq_u8_s8(z1));
3500            g += 1;
3501        }
3502    }
3503}
3504
3505#[cfg(target_arch = "x86_64")]
3506#[target_feature(enable = "avx2")]
3507unsafe fn vbit_fill4_avx2(data: &[u8], buf: &mut [u8]) {
3508    // SAFETY: see vbit_fill4_neon.
3509    unsafe {
3510        use core::arch::x86_64::*;
3511        let n = buf.len();
3512        let mask = _mm_set1_epi8(0x0F);
3513        let seven = _mm256_set1_epi8(7);
3514        let mut g = 0usize;
3515        while g * 32 + 32 <= n {
3516            let b = _mm_loadu_si128(data.as_ptr().add(g * 16) as *const __m128i);
3517            let hi = _mm_and_si128(_mm_srli_epi16::<4>(b), mask);
3518            let lo = _mm_and_si128(b, mask);
3519            let z = _mm256_sub_epi8(
3520                _mm256_set_m128i(_mm_unpackhi_epi8(hi, lo), _mm_unpacklo_epi8(hi, lo)),
3521                seven,
3522            );
3523            _mm256_storeu_si256(buf.as_mut_ptr().add(g * 32) as *mut __m256i, z);
3524            g += 1;
3525        }
3526    }
3527}
3528
3529/// Unpack 8 MSB-first B-bit values from exactly B bytes (fixed shifts —
3530/// no serial bit-buffer, auto-vectorizable). Every 32-value group starts
3531/// byte-aligned (32·B/8 is integral for B∈3..8), so groups decompose
3532/// into 4 such blocks.
3533#[inline(always)]
3534fn unpack8<const B: usize>(data: &[u8]) -> [i32; 8] {
3535    let mut acc = 0u64;
3536    for i in 0..B {
3537        acc = (acc << 8) | data[i] as u64;
3538    }
3539    let mask = (1u64 << B) - 1;
3540    let mut out = [0i32; 8];
3541    for (k, o) in out.iter_mut().enumerate() {
3542        *o = ((acc >> ((7 - k) * B)) & mask) as i32;
3543    }
3544    out
3545}
3546
3547/// Fused vbit matvec straight from the mapped bytes (spec §3, P13
3548/// FIG.3): [u8 bits: rows][f16 scales: rows·cols/32][bit-packed rows,
3549/// MSB-first, byte-padded]. Row data offsets are precomputed at load
3550/// (`vbit_row_offsets`) — the per-call prefix scan was O(rows) pure
3551/// overhead on every matvec.
3552#[allow(clippy::too_many_arguments)]
3553fn vbitmatvec(
3554    bytes: &[u8],
3555    offsets: &[usize],
3556    x: &[f32],
3557    rows: usize,
3558    cols: usize,
3559    out: &mut [f32],
3560    pool: Option<&Pool>,
3561) {
3562    debug_assert_eq!(out.len(), rows);
3563    debug_assert_eq!(offsets.len(), rows + 1);
3564
3565    // SDOT path: unpack the row to centered i8 once, then per-group
3566    // int8 dot against the quantized activations — same A8W8 contract
3567    // as q8 (bounded noise; CMF_SDOT=0 keeps the exact scalar path).
3568    if a8w8_enabled() {
3569        let act = split_act(x);
3570        let out_addr = SendMut(out.as_mut_ptr());
3571        let run = move |start: usize, end: usize| {
3572            vbit_range_a8w8(bytes, offsets, x, &act, rows, cols, out_addr, start, end)
3573        };
3574        dispatch_rows(pool, rows, &run);
3575        return;
3576    }
3577
3578    let out_addr = SendMut(out.as_mut_ptr());
3579    let run = move |start: usize, end: usize| {
3580        vbit_range_f32(bytes, offsets, x, rows, cols, out_addr, start, end)
3581    };
3582    dispatch_rows(pool, rows, &run);
3583}
3584
3585/// One vbit row range via the A8W8 int8 path — kernel body of
3586/// `vbitmatvec`, extracted so multi-matrix jobs can drive it for
3587/// several tensors in one dispatch (b=8 rows go exact f32).
3588#[allow(clippy::too_many_arguments)]
3589fn vbit_range_a8w8(
3590    bytes: &[u8],
3591    offsets: &[usize],
3592    x: &[f32],
3593    act: &SplitAct,
3594    rows: usize,
3595    cols: usize,
3596    out: SendMut,
3597    start: usize,
3598    end: usize,
3599) {
3600    let ng = cols / GROUP_SIZE;
3601    let bits = &bytes[..rows];
3602    let sc_off = rows;
3603    let row_dot = |r: usize| -> f32 {
3604        let b = bits[r] as usize;
3605        let l = (1i32 << (b - 1)) - 1;
3606        let mask = (1u64 << b) - 1;
3607        let data = &bytes[offsets[r]..offsets[r + 1]];
3608        if b == 8 {
3609            // u−L reaches 128 → does not fit i8; exact f32 path.
3610            let (mut acc, mut nbits, mut idx) = (0u64, 0usize, 0usize);
3611            let mut dot = 0f32;
3612            for g in 0..ng {
3613                let so = (r * ng + g) * 2;
3614                let sgf = f16_to_f32(u16::from_le_bytes([
3615                    bytes[sc_off + so],
3616                    bytes[sc_off + so + 1],
3617                ]));
3618                let xg = &x[g * GROUP_SIZE..(g + 1) * GROUP_SIZE];
3619                let mut gd = 0f32;
3620                for &xv in xg.iter() {
3621                    if nbits < 8 {
3622                        acc = (acc << 8) | data[idx] as u64;
3623                        idx += 1;
3624                        nbits += 8;
3625                    }
3626                    let u = ((acc >> (nbits - 8)) & 0xFF) as i32;
3627                    nbits -= 8;
3628                    gd += (u - l) as f32 * xv;
3629                }
3630                dot += gd * sgf;
3631            }
3632            return dot;
3633        }
3634        // Per-worker scratch: this closure runs for every row of the
3635        // tensor (lm_head ≈ 150k rows/token) — a heap allocation per
3636        // row was measurable pure overhead.
3637        thread_local! {
3638            static VBIT_SCRATCH: std::cell::RefCell<Vec<u8>> =
3639                const { std::cell::RefCell::new(Vec::new()) };
3640        }
3641        #[inline(always)]
3642        fn fill<const B: usize>(data: &[u8], l: i32, buf: &mut [u8]) {
3643            for (blk, chunk) in buf.chunks_exact_mut(8).enumerate() {
3644                let u = unpack8::<B>(&data[blk * B..]);
3645                for k in 0..8 {
3646                    chunk[k] = (u[k] - l) as i8 as u8;
3647                }
3648            }
3649        }
3650        let _ = mask;
3651        VBIT_SCRATCH.with(|scratch| {
3652            let mut buf = scratch.borrow_mut();
3653            buf.resize(cols, 0);
3654            match b {
3655                3 => fill::<3>(data, l, &mut buf),
3656                4 => vbit_fill4(data, &mut buf),
3657                5 => fill::<5>(data, l, &mut buf),
3658                6 => fill::<6>(data, l, &mut buf),
3659                _ => unreachable!(),
3660            }
3661            let mut dot = 0f32;
3662            for g in 0..ng {
3663                let so = (r * ng + g) * 2;
3664                let s = f16_to_f32(u16::from_le_bytes([
3665                    bytes[sc_off + so],
3666                    bytes[sc_off + so + 1],
3667                ]));
3668                let d = dot_i8_i8(
3669                    &buf[g * GROUP_SIZE..(g + 1) * GROUP_SIZE],
3670                    &act.xq[g * GROUP_SIZE..(g + 1) * GROUP_SIZE],
3671                ) as f32
3672                    * act.sx;
3673                dot += d * s;
3674            }
3675            for &(j, xv) in &act.outliers {
3676                let so = (r * ng + j / GROUP_SIZE) * 2;
3677                let s = f16_to_f32(u16::from_le_bytes([
3678                    bytes[sc_off + so],
3679                    bytes[sc_off + so + 1],
3680                ]));
3681                // xq is zeroed at outlier slots — add the exact term.
3682                dot += (buf[j] as i8) as f32 * s * xv;
3683            }
3684            dot
3685        })
3686    };
3687    for r in start..end {
3688        // SAFETY: disjoint row ranges per worker.
3689        unsafe { *out.at(r) = row_dot(r) };
3690    }
3691}
3692
3693/// Exact scalar vbit row range (same extraction, non-SDOT path).
3694#[allow(clippy::too_many_arguments)]
3695fn vbit_range_f32(
3696    bytes: &[u8],
3697    offsets: &[usize],
3698    x: &[f32],
3699    rows: usize,
3700    cols: usize,
3701    out: SendMut,
3702    start: usize,
3703    end: usize,
3704) {
3705    let ng = cols / GROUP_SIZE;
3706    let bits = &bytes[..rows];
3707    let sc_off = rows;
3708    // Per-bit-width specialized inner loops: the compiler unrolls the
3709    // constant shifts (the generic bit-buffer loop was branch-bound —
3710    // 5.6 vs 13.2 tok/s q4 on the 0.8B).
3711    #[inline(always)]
3712    fn dot_row<const B: usize>(
3713        data: &[u8],
3714        bytes: &[u8],
3715        sc_off: usize,
3716        r: usize,
3717        ng: usize,
3718        x: &[f32],
3719    ) -> f32 {
3720        let l = ((1i32 << (B - 1)) - 1) as f32;
3721        let gbytes = GROUP_SIZE * B / 8;
3722        let mut dot = 0f32;
3723        for g in 0..ng {
3724            let so = (r * ng + g) * 2;
3725            let s = f16_to_f32(u16::from_le_bytes([
3726                bytes[sc_off + so],
3727                bytes[sc_off + so + 1],
3728            ]));
3729            let xg = &x[g * GROUP_SIZE..(g + 1) * GROUP_SIZE];
3730            let gd0 = &data[g * gbytes..(g + 1) * gbytes];
3731            let mut gd = 0f32;
3732            for blk in 0..GROUP_SIZE / 8 {
3733                let u = unpack8::<B>(&gd0[blk * B..]);
3734                let xb = &xg[blk * 8..blk * 8 + 8];
3735                for k in 0..8 {
3736                    gd += (u[k] as f32 - l) * xb[k];
3737                }
3738            }
3739            dot += gd * s;
3740        }
3741        dot
3742    }
3743    for r in start..end {
3744        let data = &bytes[offsets[r]..offsets[r + 1]];
3745        let v = match bits[r] {
3746            3 => dot_row::<3>(data, bytes, sc_off, r, ng, x),
3747            4 => dot_row::<4>(data, bytes, sc_off, r, ng, x),
3748            5 => dot_row::<5>(data, bytes, sc_off, r, ng, x),
3749            6 => dot_row::<6>(data, bytes, sc_off, r, ng, x),
3750            8 => dot_row::<8>(data, bytes, sc_off, r, ng, x),
3751            b => unreachable!("vbit bit-width {b} (validated at load)"),
3752        };
3753        // SAFETY: disjoint row ranges per worker.
3754        unsafe { *out.at(r) = v };
3755    }
3756}
3757
3758/// Fused two-input vbit matvec: each row is unpacked from the mmap ONCE
3759/// and dotted against BOTH activations (MTP verify / pair prefill used
3760/// to run two full matvecs — double weight traffic and double unpack).
3761/// Per-input math is identical to `vbitmatvec` → same accuracy contract.
3762#[allow(clippy::too_many_arguments)]
3763fn vbitmatvec2(
3764    bytes: &[u8],
3765    offsets: &[usize],
3766    x1: &[f32],
3767    x2: &[f32],
3768    rows: usize,
3769    cols: usize,
3770    o1: &mut [f32],
3771    o2: &mut [f32],
3772    pool: Option<&Pool>,
3773) {
3774    debug_assert_eq!(o1.len(), rows);
3775    debug_assert_eq!(o2.len(), rows);
3776
3777    if a8w8_enabled() {
3778        let a1 = split_act(x1);
3779        let a2 = split_act(x2);
3780        let p1 = SendMut(o1.as_mut_ptr());
3781        let p2 = SendMut(o2.as_mut_ptr());
3782        let run = move |start: usize, end: usize| {
3783            vbit_range2_a8w8(
3784                bytes, offsets, x1, x2, &a1, &a2, rows, cols, p1, p2, start, end,
3785            )
3786        };
3787        dispatch_rows(pool, rows, &run);
3788        return;
3789    }
3790
3791    let p1 = SendMut(o1.as_mut_ptr());
3792    let p2 = SendMut(o2.as_mut_ptr());
3793    let run = move |start: usize, end: usize| {
3794        vbit_range2_f32(bytes, offsets, x1, x2, rows, cols, p1, p2, start, end)
3795    };
3796    dispatch_rows(pool, rows, &run);
3797}
3798
3799/// Two-input vbit row range via the A8W8 int8 path — kernel body of
3800/// `vbitmatvec2`, extracted for pair multi-matrix jobs (b=8 rows go
3801/// exact f32 for both lanes, bits streamed once).
3802#[allow(clippy::too_many_arguments)]
3803fn vbit_range2_a8w8(
3804    bytes: &[u8],
3805    offsets: &[usize],
3806    x1: &[f32],
3807    x2: &[f32],
3808    a1: &SplitAct,
3809    a2: &SplitAct,
3810    rows: usize,
3811    cols: usize,
3812    p1: SendMut,
3813    p2: SendMut,
3814    start: usize,
3815    end: usize,
3816) {
3817    let ng = cols / GROUP_SIZE;
3818    let bits = &bytes[..rows];
3819    let sc_off = rows;
3820    let row_dots = |r: usize| -> (f32, f32) {
3821        let b = bits[r] as usize;
3822        let l = (1i32 << (b - 1)) - 1;
3823        let data = &bytes[offsets[r]..offsets[r + 1]];
3824        if b == 8 {
3825            // u−L reaches 128 → does not fit i8; exact f32 path,
3826            // bits still streamed once for both lanes.
3827            let (mut acc, mut nbits, mut idx) = (0u64, 0usize, 0usize);
3828            let (mut d1, mut d2) = (0f32, 0f32);
3829            for g in 0..ng {
3830                let so = (r * ng + g) * 2;
3831                let sgf = f16_to_f32(u16::from_le_bytes([
3832                    bytes[sc_off + so],
3833                    bytes[sc_off + so + 1],
3834                ]));
3835                let (mut g1, mut g2) = (0f32, 0f32);
3836                for k in 0..GROUP_SIZE {
3837                    if nbits < 8 {
3838                        acc = (acc << 8) | data[idx] as u64;
3839                        idx += 1;
3840                        nbits += 8;
3841                    }
3842                    let u = ((acc >> (nbits - 8)) & 0xFF) as i32;
3843                    nbits -= 8;
3844                    let w = (u - l) as f32;
3845                    g1 += w * x1[g * GROUP_SIZE + k];
3846                    g2 += w * x2[g * GROUP_SIZE + k];
3847                }
3848                d1 += g1 * sgf;
3849                d2 += g2 * sgf;
3850            }
3851            return (d1, d2);
3852        }
3853        thread_local! {
3854            static VBIT_SCRATCH2: std::cell::RefCell<Vec<u8>> =
3855                const { std::cell::RefCell::new(Vec::new()) };
3856        }
3857        #[inline(always)]
3858        fn fill<const B: usize>(data: &[u8], l: i32, buf: &mut [u8]) {
3859            for (blk, chunk) in buf.chunks_exact_mut(8).enumerate() {
3860                let u = unpack8::<B>(&data[blk * B..]);
3861                for k in 0..8 {
3862                    chunk[k] = (u[k] - l) as i8 as u8;
3863                }
3864            }
3865        }
3866        VBIT_SCRATCH2.with(|scratch| {
3867            let mut buf = scratch.borrow_mut();
3868            buf.resize(cols, 0);
3869            match b {
3870                3 => fill::<3>(data, l, &mut buf),
3871                4 => vbit_fill4(data, &mut buf),
3872                5 => fill::<5>(data, l, &mut buf),
3873                6 => fill::<6>(data, l, &mut buf),
3874                _ => unreachable!(),
3875            }
3876            let (mut d1, mut d2) = (0f32, 0f32);
3877            for g in 0..ng {
3878                let so = (r * ng + g) * 2;
3879                let s = f16_to_f32(u16::from_le_bytes([
3880                    bytes[sc_off + so],
3881                    bytes[sc_off + so + 1],
3882                ]));
3883                let wg = &buf[g * GROUP_SIZE..(g + 1) * GROUP_SIZE];
3884                let v1 = dot_i8_i8(wg, &a1.xq[g * GROUP_SIZE..(g + 1) * GROUP_SIZE]) as f32 * a1.sx;
3885                let v2 = dot_i8_i8(wg, &a2.xq[g * GROUP_SIZE..(g + 1) * GROUP_SIZE]) as f32 * a2.sx;
3886                d1 += v1 * s;
3887                d2 += v2 * s;
3888            }
3889            for &(j, xv) in &a1.outliers {
3890                let so = (r * ng + j / GROUP_SIZE) * 2;
3891                let s = f16_to_f32(u16::from_le_bytes([
3892                    bytes[sc_off + so],
3893                    bytes[sc_off + so + 1],
3894                ]));
3895                d1 += (buf[j] as i8) as f32 * s * xv;
3896            }
3897            for &(j, xv) in &a2.outliers {
3898                let so = (r * ng + j / GROUP_SIZE) * 2;
3899                let s = f16_to_f32(u16::from_le_bytes([
3900                    bytes[sc_off + so],
3901                    bytes[sc_off + so + 1],
3902                ]));
3903                d2 += (buf[j] as i8) as f32 * s * xv;
3904            }
3905            (d1, d2)
3906        })
3907    };
3908    for r in start..end {
3909        let (v1, v2) = row_dots(r);
3910        // SAFETY: disjoint row ranges per worker.
3911        unsafe {
3912            *p1.at(r) = v1;
3913            *p2.at(r) = v2;
3914        }
3915    }
3916}
3917
3918/// Two-input exact scalar vbit row range (same extraction) —
3919/// per-bit-width specialized, two accumulators per row; per-lane
3920/// accumulation order matches `vbitmatvec` exactly.
3921#[allow(clippy::too_many_arguments)]
3922fn vbit_range2_f32(
3923    bytes: &[u8],
3924    offsets: &[usize],
3925    x1: &[f32],
3926    x2: &[f32],
3927    rows: usize,
3928    cols: usize,
3929    p1: SendMut,
3930    p2: SendMut,
3931    start: usize,
3932    end: usize,
3933) {
3934    let ng = cols / GROUP_SIZE;
3935    let bits = &bytes[..rows];
3936    let sc_off = rows;
3937    #[inline(always)]
3938    #[allow(clippy::too_many_arguments)]
3939    fn dot_row2<const B: usize>(
3940        data: &[u8],
3941        bytes: &[u8],
3942        sc_off: usize,
3943        r: usize,
3944        ng: usize,
3945        x1: &[f32],
3946        x2: &[f32],
3947    ) -> (f32, f32) {
3948        let l = ((1i32 << (B - 1)) - 1) as f32;
3949        let gbytes = GROUP_SIZE * B / 8;
3950        let (mut d1, mut d2) = (0f32, 0f32);
3951        for g in 0..ng {
3952            let so = (r * ng + g) * 2;
3953            let s = f16_to_f32(u16::from_le_bytes([
3954                bytes[sc_off + so],
3955                bytes[sc_off + so + 1],
3956            ]));
3957            let x1g = &x1[g * GROUP_SIZE..(g + 1) * GROUP_SIZE];
3958            let x2g = &x2[g * GROUP_SIZE..(g + 1) * GROUP_SIZE];
3959            let gd0 = &data[g * gbytes..(g + 1) * gbytes];
3960            let (mut g1, mut g2) = (0f32, 0f32);
3961            for blk in 0..GROUP_SIZE / 8 {
3962                let u = unpack8::<B>(&gd0[blk * B..]);
3963                for k in 0..8 {
3964                    let w = u[k] as f32 - l;
3965                    g1 += w * x1g[blk * 8 + k];
3966                    g2 += w * x2g[blk * 8 + k];
3967                }
3968            }
3969            d1 += g1 * s;
3970            d2 += g2 * s;
3971        }
3972        (d1, d2)
3973    }
3974    for r in start..end {
3975        let data = &bytes[offsets[r]..offsets[r + 1]];
3976        let (v1, v2) = match bits[r] {
3977            3 => dot_row2::<3>(data, bytes, sc_off, r, ng, x1, x2),
3978            4 => dot_row2::<4>(data, bytes, sc_off, r, ng, x1, x2),
3979            5 => dot_row2::<5>(data, bytes, sc_off, r, ng, x1, x2),
3980            6 => dot_row2::<6>(data, bytes, sc_off, r, ng, x1, x2),
3981            8 => dot_row2::<8>(data, bytes, sc_off, r, ng, x1, x2),
3982            b => unreachable!("vbit bit-width {b} (validated at load)"),
3983        };
3984        // SAFETY: disjoint row ranges per worker.
3985        unsafe {
3986            *p1.at(r) = v1;
3987            *p2.at(r) = v2;
3988        }
3989    }
3990}
3991
3992// ───────────────────── q4_tiled kernels (§4.3) ─────────────────────
3993
3994/// One q4_tiled row dot on the A8W8 int8 path: per 32-group the tile
3995/// is ONE sequential read — [f16 scale][16B nibbles] — versus the two
3996/// distant streams of the split layout. Values/order identical to the
3997/// split kernels.
3998#[inline]
3999#[allow(unreachable_code)]
4000fn dot_q4t_row_i8(bytes: &[u8], r: usize, gpr: usize, xq: &[i8]) -> f32 {
4001    #[cfg(target_arch = "aarch64")]
4002    unsafe {
4003        return dot_q4t_row_sdot(bytes, r, gpr, xq);
4004    }
4005    #[cfg(target_arch = "x86_64")]
4006    unsafe {
4007        if vnni_tiles_enabled() {
4008            return dot_q4t_row_vnni(bytes, r, gpr, xq);
4009        }
4010        return dot_q4t_row_avx2(bytes, r, gpr, xq);
4011    }
4012    let mut acc = 0f32;
4013    for gi in 0..gpr {
4014        let tile = &bytes[(r * gpr + gi) * Q4_TILE..(r * gpr + gi + 1) * Q4_TILE];
4015        let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
4016        let mut d = 0i32;
4017        for (k, &b) in tile[2..].iter().enumerate() {
4018            d += ((b & 0x0F) as i32 - 8) * xq[gi * GROUP_SIZE + k * 2] as i32
4019                + (((b >> 4) & 0x0F) as i32 - 8) * xq[gi * GROUP_SIZE + k * 2 + 1] as i32;
4020        }
4021        acc += d as f32 * s;
4022    }
4023    acc
4024}
4025
4026#[cfg(target_arch = "aarch64")]
4027#[target_feature(enable = "neon,dotprod")]
4028unsafe fn dot_q4t_row_sdot(bytes: &[u8], r: usize, gpr: usize, xq: &[i8]) -> f32 {
4029    // SAFETY: callers uphold slice-length contracts (18B tile per group,
4030    // xq.len() == gpr·GROUP_SIZE).
4031    unsafe {
4032        use core::arch::aarch64::*;
4033        use core::arch::asm;
4034        let lomask = vdupq_n_u8(0x0F);
4035        let eight = vdupq_n_s8(8);
4036        let mut acc = 0f32;
4037        for gi in 0..gpr {
4038            let t = bytes.as_ptr().add((r * gpr + gi) * Q4_TILE);
4039            let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
4040            let b = vld1q_u8(t.add(2));
4041            let lo = vandq_u8(b, lomask);
4042            let hi = vshrq_n_u8::<4>(b);
4043            let e0 = vsubq_s8(vreinterpretq_s8_u8(vzip1q_u8(lo, hi)), eight);
4044            let e1 = vsubq_s8(vreinterpretq_s8_u8(vzip2q_u8(lo, hi)), eight);
4045            let x0 = vld1q_s8(xq.as_ptr().add(gi * GROUP_SIZE));
4046            let x1 = vld1q_s8(xq.as_ptr().add(gi * GROUP_SIZE + 16));
4047            let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
4048            asm!(
4049                "sdot {a0:v}.4s, {e0:v}.16b, {x0:v}.16b",
4050                "sdot {a1:v}.4s, {e1:v}.16b, {x1:v}.16b",
4051                a0 = inout(vreg) a0, a1 = inout(vreg) a1,
4052                e0 = in(vreg) e0, x0 = in(vreg) x0, e1 = in(vreg) e1, x1 = in(vreg) x1,
4053                options(pure, nomem, nostack),
4054            );
4055            acc += vaddvq_s32(vaddq_s32(a0, a1)) as f32 * s;
4056        }
4057        acc
4058    }
4059}
4060
4061#[cfg(target_arch = "x86_64")]
4062#[target_feature(enable = "avx2")]
4063unsafe fn dot_q4t_row_avx2(bytes: &[u8], r: usize, gpr: usize, xq: &[i8]) -> f32 {
4064    // SAFETY: see dot_q4t_row_sdot.
4065    unsafe {
4066        use core::arch::x86_64::*;
4067        let lomask = _mm_set1_epi8(0x0F);
4068        let eight = _mm256_set1_epi8(8);
4069        let ones = _mm256_set1_epi16(1);
4070        let mut acc = 0f32;
4071        for gi in 0..gpr {
4072            let t = bytes.as_ptr().add((r * gpr + gi) * Q4_TILE);
4073            let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
4074            let b = _mm_loadu_si128(t.add(2) as *const __m128i);
4075            let lo = _mm_and_si128(b, lomask);
4076            let hi = _mm_and_si128(_mm_srli_epi16::<4>(b), lomask);
4077            let w = _mm256_sub_epi8(
4078                _mm256_set_m128i(_mm_unpackhi_epi8(lo, hi), _mm_unpacklo_epi8(lo, hi)),
4079                eight,
4080            );
4081            let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
4082            let p16 = _mm256_maddubs_epi16(_mm256_abs_epi8(w), _mm256_sign_epi8(x, w));
4083            let d = _mm256_madd_epi16(p16, ones);
4084            let hi128 = _mm256_extracti128_si256::<1>(d);
4085            let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
4086            let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
4087            let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
4088            acc += _mm_cvtsi128_si32(s32) as f32 * s;
4089        }
4090        acc
4091    }
4092}
4093
4094/// VNNI twin of `dot_q4t_row_avx2`: same unpack, `vpdpbusd` replaces
4095/// the maddubs+madd pair (see `dpbusd_hsum` — sums are bit-identical).
4096/// 256-bit VL encoding, so the VEX `vpsignb` stays usable.
4097#[cfg(target_arch = "x86_64")]
4098#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
4099unsafe fn dot_q4t_row_vnni(bytes: &[u8], r: usize, gpr: usize, xq: &[i8]) -> f32 {
4100    // SAFETY: see dot_q4t_row_sdot.
4101    unsafe {
4102        use core::arch::x86_64::*;
4103        let lomask = _mm_set1_epi8(0x0F);
4104        let eight = _mm256_set1_epi8(8);
4105        let mut acc = 0f32;
4106        for gi in 0..gpr {
4107            let t = bytes.as_ptr().add((r * gpr + gi) * Q4_TILE);
4108            let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
4109            let b = _mm_loadu_si128(t.add(2) as *const __m128i);
4110            let lo = _mm_and_si128(b, lomask);
4111            let hi = _mm_and_si128(_mm_srli_epi16::<4>(b), lomask);
4112            let w = _mm256_sub_epi8(
4113                _mm256_set_m128i(_mm_unpackhi_epi8(lo, hi), _mm_unpacklo_epi8(lo, hi)),
4114                eight,
4115            );
4116            let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
4117            let d = dpbusd_hsum(_mm256_abs_epi8(w), _mm256_sign_epi8(x, w));
4118            acc += d as f32 * s;
4119        }
4120        acc
4121    }
4122}
4123
4124/// One q4_tiled row against FOUR activation streams: the nibble unpack
4125/// and abs() happen once per group instead of once per (group,
4126/// activation) — the unpack is the dominant per-element cost of the
4127/// tiled format (roadmap P0 portable blocking, q4t leg).
4128#[cfg(target_arch = "x86_64")]
4129// `fma` is NOT implied by `avx2`: without it LLVM lowers _mm256_fmadd_ps
4130// to a libm call per lane — measured 2x slower than the reduction this
4131// kernel replaces. The runtime gate (`avx2_enabled`) already requires
4132// both features, so declaring it here is safe.
4133#[target_feature(enable = "avx2,fma")]
4134unsafe fn dot_q4t_row_1x4_avx2(bytes: &[u8], r: usize, gpr: usize, xs: [&[i8]; 4]) -> [f32; 4] {
4135    // SAFETY: callers uphold the 18B-tile and xq-length contracts.
4136    unsafe {
4137        use core::arch::x86_64::*;
4138        let lomask = _mm_set1_epi8(0x0F);
4139        let eight = _mm256_set1_epi8(8);
4140        let ones = _mm256_set1_epi16(1);
4141        // One f32 accumulator VECTOR per activation, reduced once at the
4142        // end. Folding each group's i32 lanes to a scalar inside the loop
4143        // costs an extracti128 + three shift/add + a movd — a cross-lane
4144        // dependency chain per (group, activation), 288 of them per row at
4145        // cols=2304. The per-group scale is what forces a float
4146        // accumulator; it does not force a horizontal sum.
4147        //
4148        // The four accumulators are NAMED, not an array: as `[__m256; 4]`
4149        // indexed by a loop variable LLVM keeps them in memory and every
4150        // group pays four 32-byte loads and stores. That alone made this
4151        // kernel 2x SLOWER than the per-group reduction it replaces
4152        // (measured on the EPYC box: 150 s vs 71 s for two 256² steps).
4153        let mut f0 = _mm256_setzero_ps();
4154        let mut f1 = _mm256_setzero_ps();
4155        let mut f2 = _mm256_setzero_ps();
4156        let mut f3 = _mm256_setzero_ps();
4157        for gi in 0..gpr {
4158            let t = bytes.as_ptr().add((r * gpr + gi) * Q4_TILE);
4159            let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
4160            let sv = _mm256_set1_ps(s);
4161            let bb = _mm_loadu_si128(t.add(2) as *const __m128i);
4162            let lo = _mm_and_si128(bb, lomask);
4163            let hi = _mm_and_si128(_mm_srli_epi16::<4>(bb), lomask);
4164            let w = _mm256_sub_epi8(
4165                _mm256_set_m128i(_mm_unpackhi_epi8(lo, hi), _mm_unpacklo_epi8(lo, hi)),
4166                eight,
4167            );
4168            let aw = _mm256_abs_epi8(w);
4169            let off = gi * GROUP_SIZE;
4170            let dot = |xq: &[i8]| {
4171                let x = _mm256_loadu_si256(xq.as_ptr().add(off) as *const __m256i);
4172                let p16 = _mm256_maddubs_epi16(aw, _mm256_sign_epi8(x, w));
4173                _mm256_cvtepi32_ps(_mm256_madd_epi16(p16, ones))
4174            };
4175            f0 = _mm256_fmadd_ps(dot(xs[0]), sv, f0);
4176            f1 = _mm256_fmadd_ps(dot(xs[1]), sv, f1);
4177            f2 = _mm256_fmadd_ps(dot(xs[2]), sv, f2);
4178            f3 = _mm256_fmadd_ps(dot(xs[3]), sv, f3);
4179        }
4180        [
4181            hsum256_ps(f0),
4182            hsum256_ps(f1),
4183            hsum256_ps(f2),
4184            hsum256_ps(f3),
4185        ]
4186    }
4187}
4188
4189/// Horizontal sum of eight f32 lanes — the one cross-lane reduction the
4190/// blocked kernels pay, once per row instead of once per group.
4191#[cfg(target_arch = "x86_64")]
4192#[target_feature(enable = "avx2")]
4193#[inline]
4194unsafe fn hsum256_ps(v: core::arch::x86_64::__m256) -> f32 {
4195    // SAFETY: pure register arithmetic on the caller's vector.
4196    unsafe {
4197        use core::arch::x86_64::*;
4198        let hi = _mm256_extractf128_ps::<1>(v);
4199        let s = _mm_add_ps(_mm256_castps256_ps128(v), hi);
4200        let s = _mm_add_ps(s, _mm_movehl_ps(s, s));
4201        let s = _mm_add_ss(s, _mm_shuffle_ps::<0x55>(s, s));
4202        _mm_cvtss_f32(s)
4203    }
4204}
4205
4206/// VNNI twin of `dot_q4t_row_1x4_avx2` (see `dpbusd_hsum`).
4207#[cfg(target_arch = "x86_64")]
4208#[target_feature(enable = "avx2,fma,avx512f,avx512bw,avx512vl,avx512vnni")]
4209unsafe fn dot_q4t_row_1x4_vnni(bytes: &[u8], r: usize, gpr: usize, xs: [&[i8]; 4]) -> [f32; 4] {
4210    // SAFETY: callers uphold the 18B-tile and xq-length contracts.
4211    unsafe {
4212        use core::arch::x86_64::*;
4213        let lomask = _mm_set1_epi8(0x0F);
4214        let eight = _mm256_set1_epi8(8);
4215        // Same shape as the AVX2 twin: accumulate in f32 vectors and pay
4216        // one cross-lane reduction per row, not per (group, activation).
4217        let mut f0 = _mm256_setzero_ps();
4218        let mut f1 = _mm256_setzero_ps();
4219        let mut f2 = _mm256_setzero_ps();
4220        let mut f3 = _mm256_setzero_ps();
4221        for gi in 0..gpr {
4222            let t = bytes.as_ptr().add((r * gpr + gi) * Q4_TILE);
4223            let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
4224            let sv = _mm256_set1_ps(s);
4225            let bb = _mm_loadu_si128(t.add(2) as *const __m128i);
4226            let lo = _mm_and_si128(bb, lomask);
4227            let hi = _mm_and_si128(_mm_srli_epi16::<4>(bb), lomask);
4228            let w = _mm256_sub_epi8(
4229                _mm256_set_m128i(_mm_unpackhi_epi8(lo, hi), _mm_unpacklo_epi8(lo, hi)),
4230                eight,
4231            );
4232            let aw = _mm256_abs_epi8(w);
4233            let off = gi * GROUP_SIZE;
4234            let dot = |xq: &[i8]| {
4235                let x = _mm256_loadu_si256(xq.as_ptr().add(off) as *const __m256i);
4236                _mm256_cvtepi32_ps(_mm256_dpbusd_epi32(
4237                    _mm256_setzero_si256(),
4238                    aw,
4239                    _mm256_sign_epi8(x, w),
4240                ))
4241            };
4242            f0 = _mm256_fmadd_ps(dot(xs[0]), sv, f0);
4243            f1 = _mm256_fmadd_ps(dot(xs[1]), sv, f1);
4244            f2 = _mm256_fmadd_ps(dot(xs[2]), sv, f2);
4245            f3 = _mm256_fmadd_ps(dot(xs[3]), sv, f3);
4246        }
4247        let acc = [
4248            hsum256_ps(f0),
4249            hsum256_ps(f1),
4250            hsum256_ps(f2),
4251            hsum256_ps(f3),
4252        ];
4253        acc
4254    }
4255}
4256
4257/// ARM twin of `dot_q4t_row_1x4_avx2`: one nibble unpack per group
4258/// serves FOUR activation streams. Per stream the group order and f32
4259/// accumulation match `dot_q4t_row_sdot` exactly — batch == matvec
4260/// bit-for-bit.
4261#[cfg(target_arch = "aarch64")]
4262#[target_feature(enable = "neon,dotprod")]
4263unsafe fn dot_q4t_row_1x4_sdot(bytes: &[u8], r: usize, gpr: usize, xs: [&[i8]; 4]) -> [f32; 4] {
4264    // SAFETY: callers uphold the 18B-tile and xq-length contracts.
4265    unsafe {
4266        use core::arch::aarch64::*;
4267        use core::arch::asm;
4268        let lomask = vdupq_n_u8(0x0F);
4269        let eight = vdupq_n_s8(8);
4270        let mut acc = [0f32; 4];
4271        for gi in 0..gpr {
4272            let t = bytes.as_ptr().add((r * gpr + gi) * Q4_TILE);
4273            let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
4274            let b = vld1q_u8(t.add(2));
4275            let lo = vandq_u8(b, lomask);
4276            let hi = vshrq_n_u8::<4>(b);
4277            let e0 = vsubq_s8(vreinterpretq_s8_u8(vzip1q_u8(lo, hi)), eight);
4278            let e1 = vsubq_s8(vreinterpretq_s8_u8(vzip2q_u8(lo, hi)), eight);
4279            for (k, xq) in xs.iter().enumerate() {
4280                let x0 = vld1q_s8(xq.as_ptr().add(gi * GROUP_SIZE));
4281                let x1 = vld1q_s8(xq.as_ptr().add(gi * GROUP_SIZE + 16));
4282                let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
4283                asm!(
4284                    "sdot {a0:v}.4s, {e0:v}.16b, {x0:v}.16b",
4285                    "sdot {a1:v}.4s, {e1:v}.16b, {x1:v}.16b",
4286                    a0 = inout(vreg) a0, a1 = inout(vreg) a1,
4287                    e0 = in(vreg) e0, x0 = in(vreg) x0, e1 = in(vreg) e1, x1 = in(vreg) x1,
4288                    options(pure, nomem, nostack),
4289                );
4290                acc[k] += vaddvq_s32(vaddq_s32(a0, a1)) as f32 * s;
4291            }
4292        }
4293        acc
4294    }
4295}
4296
4297/// Exact-term correction for A8W8 outliers on a tiled row.
4298#[inline]
4299fn q4t_outlier(bytes: &[u8], r: usize, gpr: usize, j: usize) -> (f32, f32) {
4300    let gi = j / GROUP_SIZE;
4301    let k = j % GROUP_SIZE;
4302    let tile = &bytes[(r * gpr + gi) * Q4_TILE..(r * gpr + gi + 1) * Q4_TILE];
4303    let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
4304    let byte = tile[2 + k / 2];
4305    let nib = if k & 1 == 0 { byte & 0x0F } else { byte >> 4 };
4306    ((nib as i32 - 8) as f32, s)
4307}
4308
4309/// Exact scalar q4_tiled row (CMF_SDOT=0 contract) — same pairwise
4310/// accumulation shape as `q4_range_f32`.
4311#[inline]
4312fn q4t_row_exact(bytes: &[u8], r: usize, gpr: usize, x: &[f32]) -> f32 {
4313    let mut acc = 0f32;
4314    for gi in 0..gpr {
4315        let tile = &bytes[(r * gpr + gi) * Q4_TILE..(r * gpr + gi + 1) * Q4_TILE];
4316        let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
4317        let xg = &x[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
4318        let mut ga = 0f32;
4319        for (k, &b) in tile[2..].iter().enumerate() {
4320            ga += ((b & 0x0F) as f32 - 8.0) * xg[k * 2]
4321                + (((b >> 4) & 0x0F) as f32 - 8.0) * xg[k * 2 + 1];
4322        }
4323        acc += ga * s;
4324    }
4325    acc
4326}
4327
4328/// Split view of a `q4tp` payload. The three planes are resolved once per
4329/// matvec instead of per row — `q4tp_sections` is cheap, but doing it inside
4330/// the row loop would put a division on the hot path for nothing.
4331struct Q4tpView<'a> {
4332    nib: &'a [u8],
4333    params: &'a [u8],
4334    codes: &'a [u8],
4335    stride: usize,
4336    /// q2tp reads the ladder with rung 0 = exact zero.
4337    zero_rung: bool,
4338}
4339
4340impl<'a> Q4tpView<'a> {
4341    fn new(bytes: &'a [u8], rows: usize, cols: usize) -> Self {
4342        let (params_off, codes_off, stride) = q4tp_sections(rows, cols);
4343        Self {
4344            nib: &bytes[..params_off],
4345            params: &bytes[params_off..codes_off],
4346            codes: &bytes[codes_off..],
4347            stride,
4348            zero_rung: false,
4349        }
4350    }
4351
4352    /// The q2tp view: identical params/codes planes, 8 B weight chunks.
4353    fn new_q2(bytes: &'a [u8], rows: usize, cols: usize) -> Self {
4354        let (params_off, codes_off, stride) = q2tp_sections(rows, cols);
4355        Self {
4356            nib: &bytes[..params_off],
4357            params: &bytes[params_off..codes_off],
4358            codes: &bytes[codes_off..],
4359            stride,
4360            zero_rung: true,
4361        }
4362    }
4363
4364    /// Expand row `r`'s per-tile scales into `out` (length `gpr`).
4365    ///
4366    /// Doing this once per row — rather than decoding a 5-bit code inside the
4367    /// tile loop — is what makes the format free at runtime. Random access to
4368    /// a packed 5-bit field costs a division, two bounds checks and a branch;
4369    /// the tile's actual work is two `sdot`s, so per-tile decoding dominated
4370    /// the kernel and cost 5x (measured: 1.4 vs 6.9 tok/s on Nanbeige-3B).
4371    /// Walking the plane sequentially with a bit accumulator is ~3 ops.
4372    /// Eight 5-bit codes are exactly five bytes, so a whole group of
4373    /// eight decodes from one little-endian word at fixed shifts. The
4374    /// bit-accumulator this replaces carried a data-dependent `while
4375    /// have < 5` refill whose branch sat in the innermost loop of every
4376    /// q4tp row; a decode profile put this function above the dot
4377    /// products it feeds. Same bitstream, same codes — just no branch
4378    /// and eight independent extractions.
4379    #[inline]
4380    fn scales_into(&self, r: usize, gpr: usize, out: &mut [f32]) {
4381        let tab = if self.zero_rung {
4382            q2tp_ladder(self.params, r)
4383        } else {
4384            q4tp_ladder(self.params, r)
4385        };
4386        let codes = &self.codes[r * self.stride..(r + 1) * self.stride];
4387        let out = &mut out[..gpr];
4388        let mut chunks = out.chunks_exact_mut(8);
4389        let mut ci = 0usize;
4390        for c in &mut chunks {
4391            let w = u64::from(codes[ci])
4392                | u64::from(codes[ci + 1]) << 8
4393                | u64::from(codes[ci + 2]) << 16
4394                | u64::from(codes[ci + 3]) << 24
4395                | u64::from(codes[ci + 4]) << 32;
4396            for (k, o) in c.iter_mut().enumerate() {
4397                *o = tab[((w >> (5 * k)) & 31) as usize];
4398            }
4399            ci += 5;
4400        }
4401        // Fewer than eight codes left: the shared total accessor, which
4402        // tolerates a 5-bit field whose spill byte is past the stride.
4403        let tail = &codes[ci..];
4404        for (k, o) in chunks.into_remainder().iter_mut().enumerate() {
4405            *o = tab[q4tp_code(tail, k)];
4406        }
4407    }
4408}
4409
4410#[inline]
4411fn dot_q4tp_row_i8(nib: &[u8], r: usize, gpr: usize, xq: &[i8], scales: &[f32]) -> f32 {
4412    #[cfg(target_arch = "aarch64")]
4413    unsafe {
4414        return dot_q4tp_row_sdot(nib, r, gpr, xq, scales);
4415    }
4416    #[cfg(target_arch = "x86_64")]
4417    unsafe {
4418        if vnni_tiles_enabled() {
4419            return dot_q4tp_row_vnni(nib, r, gpr, xq, scales);
4420        }
4421        return dot_q4tp_row_avx2(nib, r, gpr, xq, scales);
4422    }
4423    #[allow(unreachable_code)]
4424    {
4425        let mut acc = 0f32;
4426        for gi in 0..gpr {
4427            let tile = &nib[(r * gpr + gi) * Q4TP_NIB..(r * gpr + gi + 1) * Q4TP_NIB];
4428            let s = scales[gi];
4429            let mut d = 0i32;
4430            for (k, &b) in tile.iter().enumerate() {
4431                d += ((b & 0x0F) as i32 - 8) * xq[gi * GROUP_SIZE + k * 2] as i32
4432                    + (((b >> 4) & 0x0F) as i32 - 8) * xq[gi * GROUP_SIZE + k * 2 + 1] as i32;
4433            }
4434            acc += d as f32 * s;
4435        }
4436        acc
4437    }
4438}
4439
4440/// q4tp twin of `dot_q4t_row_sdot`: identical nibble math, but the tile
4441/// stride is 16 B (no inline scale) and the scale is a ladder lookup.
4442#[cfg(target_arch = "aarch64")]
4443#[target_feature(enable = "neon,dotprod")]
4444unsafe fn dot_q4tp_row_sdot(nib: &[u8], r: usize, gpr: usize, xq: &[i8], scales: &[f32]) -> f32 {
4445    // SAFETY: callers uphold slice-length contracts (16B tile per group,
4446    // xq.len() == gpr·GROUP_SIZE, codes covering gpr 5-bit fields).
4447    unsafe {
4448        use core::arch::aarch64::*;
4449        use core::arch::asm;
4450        let lomask = vdupq_n_u8(0x0F);
4451        let eight = vdupq_n_s8(8);
4452        let mut acc = 0f32;
4453        for gi in 0..gpr {
4454            let t = nib.as_ptr().add((r * gpr + gi) * Q4TP_NIB);
4455            let s = *scales.get_unchecked(gi);
4456            let b = vld1q_u8(t);
4457            let lo = vandq_u8(b, lomask);
4458            let hi = vshrq_n_u8::<4>(b);
4459            let e0 = vsubq_s8(vreinterpretq_s8_u8(vzip1q_u8(lo, hi)), eight);
4460            let e1 = vsubq_s8(vreinterpretq_s8_u8(vzip2q_u8(lo, hi)), eight);
4461            let x0 = vld1q_s8(xq.as_ptr().add(gi * GROUP_SIZE));
4462            let x1 = vld1q_s8(xq.as_ptr().add(gi * GROUP_SIZE + 16));
4463            let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
4464            asm!(
4465                "sdot {a0:v}.4s, {e0:v}.16b, {x0:v}.16b",
4466                "sdot {a1:v}.4s, {e1:v}.16b, {x1:v}.16b",
4467                a0 = inout(vreg) a0, a1 = inout(vreg) a1,
4468                e0 = in(vreg) e0, x0 = in(vreg) x0, e1 = in(vreg) e1, x1 = in(vreg) x1,
4469                options(pure, nomem, nostack),
4470            );
4471            acc += vaddvq_s32(vaddq_s32(a0, a1)) as f32 * s;
4472        }
4473        acc
4474    }
4475}
4476
4477#[cfg(target_arch = "x86_64")]
4478#[target_feature(enable = "avx2")]
4479unsafe fn dot_q4tp_row_avx2(nib: &[u8], r: usize, gpr: usize, xq: &[i8], scales: &[f32]) -> f32 {
4480    // SAFETY: see dot_q4tp_row_sdot.
4481    unsafe {
4482        use core::arch::x86_64::*;
4483        let lomask = _mm_set1_epi8(0x0F);
4484        let eight = _mm256_set1_epi8(8);
4485        let ones = _mm256_set1_epi16(1);
4486        let mut acc = 0f32;
4487        for gi in 0..gpr {
4488            let t = nib.as_ptr().add((r * gpr + gi) * Q4TP_NIB);
4489            let s = *scales.get_unchecked(gi);
4490            let b = _mm_loadu_si128(t as *const __m128i);
4491            let lo = _mm_and_si128(b, lomask);
4492            let hi = _mm_and_si128(_mm_srli_epi16::<4>(b), lomask);
4493            let w = _mm256_sub_epi8(
4494                _mm256_set_m128i(_mm_unpackhi_epi8(lo, hi), _mm_unpacklo_epi8(lo, hi)),
4495                eight,
4496            );
4497            let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
4498            let p16 = _mm256_maddubs_epi16(_mm256_abs_epi8(w), _mm256_sign_epi8(x, w));
4499            let d = _mm256_madd_epi16(p16, ones);
4500            let hi128 = _mm256_extracti128_si256::<1>(d);
4501            let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
4502            let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
4503            let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
4504            acc += _mm_cvtsi128_si32(s32) as f32 * s;
4505        }
4506        acc
4507    }
4508}
4509
4510/// VNNI twin of `dot_q4tp_row_avx2` (see `dot_q4t_row_vnni` for why the
4511/// 256-bit VL encoding is the one to use here).
4512#[cfg(target_arch = "x86_64")]
4513#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
4514unsafe fn dot_q4tp_row_vnni(nib: &[u8], r: usize, gpr: usize, xq: &[i8], scales: &[f32]) -> f32 {
4515    // SAFETY: see dot_q4tp_row_sdot.
4516    unsafe {
4517        use core::arch::x86_64::*;
4518        let lomask = _mm_set1_epi8(0x0F);
4519        let eight = _mm256_set1_epi8(8);
4520        let mut acc = 0f32;
4521        for gi in 0..gpr {
4522            let t = nib.as_ptr().add((r * gpr + gi) * Q4TP_NIB);
4523            let s = *scales.get_unchecked(gi);
4524            let b = _mm_loadu_si128(t as *const __m128i);
4525            let lo = _mm_and_si128(b, lomask);
4526            let hi = _mm_and_si128(_mm_srli_epi16::<4>(b), lomask);
4527            let w = _mm256_sub_epi8(
4528                _mm256_set_m128i(_mm_unpackhi_epi8(lo, hi), _mm_unpacklo_epi8(lo, hi)),
4529                eight,
4530            );
4531            let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
4532            acc += dpbusd_hsum(_mm256_abs_epi8(w), _mm256_sign_epi8(x, w)) as f32 * s;
4533        }
4534        acc
4535    }
4536}
4537
4538/// Exact scalar q4tp row — the `CMF_SDOT=0` contract, same pairwise
4539/// accumulation shape as `q4t_row_exact`.
4540#[inline]
4541fn q4tp_row_exact(nib: &[u8], r: usize, gpr: usize, x: &[f32], scales: &[f32]) -> f32 {
4542    let mut acc = 0f32;
4543    for gi in 0..gpr {
4544        let tile = &nib[(r * gpr + gi) * Q4TP_NIB..(r * gpr + gi + 1) * Q4TP_NIB];
4545        let s = scales[gi];
4546        let xg = &x[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
4547        let mut ga = 0f32;
4548        for (k, &b) in tile.iter().enumerate() {
4549            ga += ((b & 0x0F) as f32 - 8.0) * xg[k * 2]
4550                + (((b >> 4) & 0x0F) as f32 - 8.0) * xg[k * 2 + 1];
4551        }
4552        acc += ga * s;
4553    }
4554    acc
4555}
4556
4557/// Single weight of a q4tp tensor — the a8w8 outlier path, which restores
4558/// activation outliers at full precision after the int8 pass.
4559#[inline]
4560fn q4tp_outlier(nib: &[u8], r: usize, gpr: usize, j: usize, scales: &[f32]) -> (f32, f32) {
4561    let (gi, k) = (j / GROUP_SIZE, j % GROUP_SIZE);
4562    let byte = nib[(r * gpr + gi) * Q4TP_NIB + k / 2];
4563    let n = if k & 1 == 0 { byte & 0x0F } else { byte >> 4 };
4564    ((n as i32 - 8) as f32, scales[gi])
4565}
4566
4567/// Fused q4tp matvec (dispatch mirrors `q4t_matvec`).
4568fn q4tp_matvec(
4569    bytes: &[u8],
4570    x: &[f32],
4571    rows: usize,
4572    cols: usize,
4573    out: &mut [f32],
4574    pool: Option<&Pool>,
4575) {
4576    debug_assert_eq!(out.len(), rows);
4577    let gpr = cols / GROUP_SIZE;
4578    let v = Q4tpView::new(bytes, rows, cols);
4579    let out_addr = SendMut(out.as_mut_ptr());
4580    if a8w8_enabled() {
4581        let act = split_act(x);
4582        let run = |start: usize, end: usize| {
4583            // One scratch row of scales per worker — borrowed, not minted.
4584            with_krow(gpr, |sc| {
4585                for r in start..end {
4586                    v.scales_into(r, gpr, sc);
4587                    let mut acc = dot_q4tp_row_i8(v.nib, r, gpr, &act.xq, sc) * act.sx;
4588                    for &(j, xv) in &act.outliers {
4589                        let (w, s) = q4tp_outlier(v.nib, r, gpr, j, sc);
4590                        acc += w * s * xv;
4591                    }
4592                    // SAFETY: disjoint row ranges per worker.
4593                    unsafe { *out_addr.at(r) = acc };
4594                }
4595            })
4596        };
4597        dispatch_rows(pool, rows, &run);
4598        return;
4599    }
4600    let run = |start: usize, end: usize| {
4601        with_krow(gpr, |sc| {
4602            for r in start..end {
4603                v.scales_into(r, gpr, sc);
4604                // SAFETY: disjoint row ranges per worker.
4605                unsafe { *out_addr.at(r) = q4tp_row_exact(v.nib, r, gpr, x, sc) };
4606            }
4607        })
4608    };
4609    dispatch_rows(pool, rows, &run);
4610}
4611
4612/// Fused two-input q4tp matvec — the SwiGLU gate/up pair. Weights and the
4613/// row ladder are read once and spent on both activation streams.
4614#[allow(clippy::too_many_arguments)]
4615fn q4tp_matvec2(
4616    bytes: &[u8],
4617    x1: &[f32],
4618    x2: &[f32],
4619    rows: usize,
4620    cols: usize,
4621    o1: &mut [f32],
4622    o2: &mut [f32],
4623    pool: Option<&Pool>,
4624) {
4625    let gpr = cols / GROUP_SIZE;
4626    let v = Q4tpView::new(bytes, rows, cols);
4627    let (p1, p2) = (SendMut(o1.as_mut_ptr()), SendMut(o2.as_mut_ptr()));
4628    let run = |start: usize, end: usize| {
4629        let mut sc = vec![0f32; gpr];
4630        for r in start..end {
4631            v.scales_into(r, gpr, &mut sc);
4632            // SAFETY: disjoint row ranges per worker.
4633            unsafe {
4634                *p1.at(r) = q4tp_row_exact(v.nib, r, gpr, x1, &sc);
4635                *p2.at(r) = q4tp_row_exact(v.nib, r, gpr, x2, &sc);
4636            }
4637        }
4638    };
4639    dispatch_rows(pool, rows, &run);
4640}
4641
4642/// One q2tp outlier weight at column `j` of row `r`: the 2-bit code and
4643/// its group scale, mirrored on `q4tp_outlier`.
4644#[inline]
4645fn q2tp_outlier(chunks: &[u8], r: usize, gpr: usize, j: usize, scales: &[f32]) -> (f32, f32) {
4646    let (gi, k) = (j / GROUP_SIZE, j % GROUP_SIZE);
4647    let byte = chunks[(r * gpr + gi) * Q2TP_CHUNK + k / 4];
4648    let c = (byte >> (2 * (k % 4))) & 3;
4649    (c as f32 - 1.5, scales[gi])
4650}
4651
4652/// Integer dot of one q2tp row against pre-quantized activations:
4653/// Σ_g s_g · (Σ c·xq − 1.5·Σ xq). The half-integer grid (c − 1.5)
4654/// becomes exact integer math through the group sums — the same trick
4655/// every a8w8 kernel in this file rides. The codes decode into a
4656/// 32-byte scratch in natural order and the dot itself is the shared
4657/// SDOT primitive; elsewhere a scalar integer loop.
4658#[inline]
4659fn dot_q2tp_row_i8(
4660    chunks: &[u8],
4661    r: usize,
4662    gpr: usize,
4663    xq: &[i8],
4664    gsum: &[i32],
4665    scales: &[f32],
4666) -> f32 {
4667    let mut acc = 0f32;
4668    let base = r * gpr * Q2TP_CHUNK;
4669    #[cfg(not(target_arch = "aarch64"))]
4670    let mut codes = [0i8; GROUP_SIZE];
4671    for gi in 0..gpr {
4672        let ch = &chunks[base + gi * Q2TP_CHUNK..base + (gi + 1) * Q2TP_CHUNK];
4673        let xg = &xq[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
4674        #[cfg(target_arch = "aarch64")]
4675        // NEON: the byte's four 2-bit fields land in four lane vectors
4676        // (shift+mask), vld4 de-interleaves xq to match (xj[k] =
4677        // xq[4k+j]), widening MACs accumulate exactly in i32. A scalar
4678        // decode here cost as much as the dot it fed — the profile put
4679        // it at the top of the whole W2 decode.
4680        let dot = unsafe {
4681            use core::arch::aarch64::*;
4682            let b = vld1_u8(ch.as_ptr());
4683            let three = vdup_n_u8(3);
4684            let c0 = vreinterpret_s8_u8(vand_u8(b, three));
4685            let c1 = vreinterpret_s8_u8(vand_u8(vshr_n_u8(b, 2), three));
4686            let c2 = vreinterpret_s8_u8(vand_u8(vshr_n_u8(b, 4), three));
4687            let c3 = vreinterpret_s8_u8(vand_u8(vshr_n_u8(b, 6), three));
4688            let x4 = vld4_s8(xg.as_ptr());
4689            let mut acc4 = vdupq_n_s32(0);
4690            acc4 = vpadalq_s16(acc4, vmull_s8(c0, x4.0));
4691            acc4 = vpadalq_s16(acc4, vmull_s8(c1, x4.1));
4692            acc4 = vpadalq_s16(acc4, vmull_s8(c2, x4.2));
4693            acc4 = vpadalq_s16(acc4, vmull_s8(c3, x4.3));
4694            vaddvq_s32(acc4)
4695        };
4696        #[cfg(not(target_arch = "aarch64"))]
4697        let dot: i32 = {
4698            for (k, &b) in ch.iter().enumerate() {
4699                codes[k * 4] = (b & 3) as i8;
4700                codes[k * 4 + 1] = ((b >> 2) & 3) as i8;
4701                codes[k * 4 + 2] = ((b >> 4) & 3) as i8;
4702                codes[k * 4 + 3] = ((b >> 6) & 3) as i8;
4703            }
4704            codes
4705                .iter()
4706                .zip(xg)
4707                .map(|(&c, &x)| c as i32 * x as i32)
4708                .sum()
4709        };
4710        acc += scales[gi] * (dot as f32 - 1.5 * gsum[gi] as f32);
4711    }
4712    acc
4713}
4714
4715/// Exact f32 dot of one q2tp row: 2-bit fields LSB-first, (c − 1.5)·s.
4716/// Scalar on purpose — the 2-bit class targets the GPU graph; the CPU
4717/// path exists for parity gates and small-machine fallback.
4718fn q2tp_row_exact(chunks: &[u8], r: usize, gpr: usize, x: &[f32], scales: &[f32]) -> f32 {
4719    let mut acc = 0f32;
4720    for gi in 0..gpr {
4721        let ch = &chunks[(r * gpr + gi) * Q2TP_CHUNK..(r * gpr + gi + 1) * Q2TP_CHUNK];
4722        let s = scales[gi];
4723        let xb = &x[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
4724        let mut g = 0f32;
4725        for (k, &b) in ch.iter().enumerate() {
4726            g += ((b & 3) as f32 - 1.5) * xb[k * 4]
4727                + (((b >> 2) & 3) as f32 - 1.5) * xb[k * 4 + 1]
4728                + (((b >> 4) & 3) as f32 - 1.5) * xb[k * 4 + 2]
4729                + (((b >> 6) & 3) as f32 - 1.5) * xb[k * 4 + 3];
4730        }
4731        acc += s * g;
4732    }
4733    acc
4734}
4735
4736fn q2tp_matvec(
4737    bytes: &[u8],
4738    x: &[f32],
4739    rows: usize,
4740    cols: usize,
4741    out: &mut [f32],
4742    pool: Option<&Pool>,
4743) {
4744    debug_assert_eq!(out.len(), rows);
4745    let gpr = cols / GROUP_SIZE;
4746    let v = Q4tpView::new_q2(bytes, rows, cols);
4747    let out_addr = SendMut(out.as_mut_ptr());
4748    // a8w8 fast path (CMF_SDOT=0 keeps the exact scalar walk): integer
4749    // code dots + group sums, exact outlier correction — the same
4750    // contract as every sibling kernel; measured 2-bit rows were the
4751    // only scalar holdout in the family.
4752    if a8w8_enabled() {
4753        let act = split_act(x);
4754        let gsum = q1_group_sums(&act.xq, gpr);
4755        let (act, gsum) = (&act, &gsum);
4756        let run = move |start: usize, end: usize| {
4757            with_krow(gpr, |sc| {
4758                for r in start..end {
4759                    v.scales_into(r, gpr, sc);
4760                    let mut acc = dot_q2tp_row_i8(v.nib, r, gpr, &act.xq, gsum, sc) * act.sx;
4761                    for &(j, xv) in &act.outliers {
4762                        let (w, s) = q2tp_outlier(v.nib, r, gpr, j, sc);
4763                        acc += w * s * xv;
4764                    }
4765                    // SAFETY: disjoint row ranges per worker.
4766                    unsafe { *out_addr.at(r) = acc };
4767                }
4768            })
4769        };
4770        dispatch_rows(pool, rows, &run);
4771        return;
4772    }
4773    let run = |start: usize, end: usize| {
4774        with_krow(gpr, |sc| {
4775            for r in start..end {
4776                v.scales_into(r, gpr, sc);
4777                // SAFETY: disjoint row ranges per worker.
4778                unsafe { *out_addr.at(r) = q2tp_row_exact(v.nib, r, gpr, x, sc) };
4779            }
4780        })
4781    };
4782    dispatch_rows(pool, rows, &run);
4783}
4784
4785/// Fused two-input q2tp matvec — the SwiGLU gate/up pair.
4786#[allow(clippy::too_many_arguments)]
4787fn q2tp_matvec2(
4788    bytes: &[u8],
4789    x1: &[f32],
4790    x2: &[f32],
4791    rows: usize,
4792    cols: usize,
4793    o1: &mut [f32],
4794    o2: &mut [f32],
4795    pool: Option<&Pool>,
4796) {
4797    let gpr = cols / GROUP_SIZE;
4798    let v = Q4tpView::new_q2(bytes, rows, cols);
4799    let (p1, p2) = (SendMut(o1.as_mut_ptr()), SendMut(o2.as_mut_ptr()));
4800    let run = |start: usize, end: usize| {
4801        let mut sc = vec![0f32; gpr];
4802        for r in start..end {
4803            v.scales_into(r, gpr, &mut sc);
4804            // SAFETY: disjoint row ranges per worker.
4805            unsafe {
4806                *p1.at(r) = q2tp_row_exact(v.nib, r, gpr, x1, &sc);
4807                *p2.at(r) = q2tp_row_exact(v.nib, r, gpr, x2, &sc);
4808            }
4809        }
4810    };
4811    dispatch_rows(pool, rows, &run);
4812}
4813
4814/// Batched q2tp matmat: scalar row kernel over every batch column. CPU
4815/// prefill only — decode rides the graph, so plain and correct beats
4816/// clever here.
4817/// Test doors into the host 2-bit kernels: the stand's heap corruption
4818/// pointed at down-shaped tensors, and the private fns need a way to be
4819/// held to a reference without a model file around them.
4820pub fn q2tp_matvec_for_test(bytes: &[u8], x: &[f32], rows: usize, cols: usize, out: &mut [f32]) {
4821    // The facade IS the reference: encoder oracles hold requant output
4822    // to the exact scalar walk. The production dispatch may take the i8
4823    // fast path, whose error scale is the ACTIVATIONS' — a different
4824    // claim than the encoder correctness these tests pin.
4825    let gpr = cols / GROUP_SIZE;
4826    let v = Q4tpView::new_q2(bytes, rows, cols);
4827    with_krow(gpr, |sc| {
4828        for r in 0..rows {
4829            v.scales_into(r, gpr, sc);
4830            out[r] = q2tp_row_exact(v.nib, r, gpr, x, sc);
4831        }
4832    });
4833}
4834
4835pub fn q2tp_matmat_for_test(
4836    bytes: &[u8],
4837    xs_all: &[f32],
4838    b: usize,
4839    rows: usize,
4840    cols: usize,
4841    out: &mut [f32],
4842) {
4843    q2tp_matmat(bytes, xs_all, b, rows, cols, out, None);
4844}
4845
4846fn q2tp_matmat(
4847    bytes: &[u8],
4848    xs_all: &[f32],
4849    b: usize,
4850    rows: usize,
4851    cols: usize,
4852    out: &mut [f32],
4853    pool: Option<&Pool>,
4854) {
4855    debug_assert_eq!(out.len(), b * rows);
4856    let gpr = cols / GROUP_SIZE;
4857    let v = Q4tpView::new_q2(bytes, rows, cols);
4858    let out_addr = SendMut(out.as_mut_ptr());
4859    let run = |start: usize, end: usize| {
4860        let mut sc = vec![0f32; gpr];
4861        for r in start..end {
4862            v.scales_into(r, gpr, &mut sc);
4863            for bi in 0..b {
4864                let x = &xs_all[bi * cols..(bi + 1) * cols];
4865                // SAFETY: disjoint row ranges per worker.
4866                unsafe { *out_addr.at(bi * rows + r) = q2tp_row_exact(v.nib, r, gpr, x, &sc) };
4867            }
4868        }
4869    };
4870    dispatch_rows(pool, rows, &run);
4871}
4872
4873/// The pre-vectorised shape, kept for A/B (`CMF_Q4TP_V1=1`): the
4874/// horizontal add lands once per group per column instead of once per
4875/// row. Same weights, same activations — only the reduction differs.
4876#[cfg(target_arch = "aarch64")]
4877#[target_feature(enable = "neon,dotprod")]
4878unsafe fn dot_q4tp_row_1x4_sdot_v1(
4879    nib: &[u8],
4880    r: usize,
4881    gpr: usize,
4882    xs: [&[i8]; 4],
4883    scales: &[f32],
4884) -> [f32; 4] {
4885    unsafe {
4886        use core::arch::aarch64::*;
4887        use core::arch::asm;
4888        let lomask = vdupq_n_u8(0x0F);
4889        let eight = vdupq_n_s8(8);
4890        let (mut f0, mut f1, mut f2, mut f3) = (0f32, 0f32, 0f32, 0f32);
4891        for gi in 0..gpr {
4892            let t = nib.as_ptr().add((r * gpr + gi) * Q4TP_NIB);
4893            let s = *scales.get_unchecked(gi);
4894            let bb = vld1q_u8(t);
4895            let lo = vandq_u8(bb, lomask);
4896            let hi = vshrq_n_u8::<4>(bb);
4897            let e0 = vsubq_s8(vreinterpretq_s8_u8(vzip1q_u8(lo, hi)), eight);
4898            let e1 = vsubq_s8(vreinterpretq_s8_u8(vzip2q_u8(lo, hi)), eight);
4899            let mut d = [0f32; 4];
4900            for (k, dk) in d.iter_mut().enumerate() {
4901                let x0 = vld1q_s8(xs[k].as_ptr().add(gi * GROUP_SIZE));
4902                let x1 = vld1q_s8(xs[k].as_ptr().add(gi * GROUP_SIZE + 16));
4903                let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
4904                asm!(
4905                    "sdot {a0:v}.4s, {e0:v}.16b, {x0:v}.16b",
4906                    "sdot {a1:v}.4s, {e1:v}.16b, {x1:v}.16b",
4907                    a0 = inout(vreg) a0, a1 = inout(vreg) a1,
4908                    e0 = in(vreg) e0, x0 = in(vreg) x0, e1 = in(vreg) e1, x1 = in(vreg) x1,
4909                    options(pure, nomem, nostack),
4910                );
4911                *dk = vaddvq_s32(vaddq_s32(a0, a1)) as f32 * s;
4912            }
4913            f0 += d[0];
4914            f1 += d[1];
4915            f2 += d[2];
4916            f3 += d[3];
4917        }
4918        [f0, f1, f2, f3]
4919    }
4920}
4921
4922/// Which q4tp batch kernel to run: 1 = the previous one, 2 = the tuned
4923/// one, 0 = decide from the CPU. An atomic rather than a `OnceLock` so a
4924/// benchmark can alternate the two inside one process, where the machine's
4925/// mood — a shared box drifts ±25% between runs — is the same for both.
4926/// What the two mean is per-architecture: on x86 the blocked AVX-512 path
4927/// against the per-column one, on ARM the two reduction shapes.
4928#[allow(dead_code)]
4929static Q4TP_ALT: std::sync::atomic::AtomicU8 = std::sync::atomic::AtomicU8::new(0);
4930
4931/// Blocking pays on x86 only with 512-bit VNNI. With AVX2 alone, four
4932/// columns sharing an unpack still measured slower than the per-column
4933/// path (23.2 ms against 19.4 on a 48-thread EPYC), because that path
4934/// already dequantizes the row once — so the blocked kernel bought a
4935/// second unpack-free pass at the price of half the vector width.
4936#[cfg(target_arch = "x86_64")]
4937fn q4tp_blocked_x86() -> bool {
4938    match Q4TP_ALT.load(std::sync::atomic::Ordering::Relaxed) {
4939        1 => false,
4940        // A forced ON still asks the CPU. The switch exists so a bench can
4941        // pick a kernel, not so it can promise instructions the machine
4942        // does not have — CI caught that as a SIGILL on a runner without
4943        // AVX-512, where the parity test had turned the path on by hand.
4944        2 => avx512vnni_enabled(),
4945        // Deliberately not cached back into the switch: both gates below
4946        // hold their own `OnceLock`, and latching their answer here would
4947        // make a test's override outlive the test that set it.
4948        _ => blocked_enabled() && avx512vnni_enabled(),
4949    }
4950}
4951
4952/// `CMF_Q4TP_V1=1` picks the old reduction shape (A/B only).
4953#[cfg(target_arch = "aarch64")]
4954#[allow(dead_code)]
4955fn q4tp_v1() -> bool {
4956    match Q4TP_ALT.load(std::sync::atomic::Ordering::Relaxed) {
4957        1 => true,
4958        2 => false,
4959        _ => {
4960            static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
4961            *ON.get_or_init(|| std::env::var("CMF_Q4TP_V1").is_ok_and(|v| v != "0"))
4962        }
4963    }
4964}
4965
4966/// Two weight rows against eight columns. The activation load is the
4967/// same for both rows, so it is paid once for twice the arithmetic, and
4968/// sixteen accumulator chains run where eight did — which is what a kernel
4969/// retiring 0.29 instructions a cycle is short of. Register pressure is
4970/// the limit: sixteen `zmm` accumulators, two weight tiles, one
4971/// activation, of thirty-two.
4972///
4973/// Four rows by four columns spends the same sixteen accumulators the
4974/// other way and measured worse — 1488 GFLOP/s against 1644 — so the
4975/// unpack, which four rows pay twice as often, costs more than the extra
4976/// sharing of one activation load buys.
4977#[cfg(target_arch = "x86_64")]
4978#[target_feature(enable = "avx512f,avx512bw,avx512vnni")]
4979unsafe fn dot_q4tp_2x8_avx512(
4980    nib: &[u8],
4981    r0: usize,
4982    gpr: usize,
4983    xs: [&[i8]; 8],
4984    sc0: &[f32],
4985    sc1: &[f32],
4986) -> [[f32; 8]; 2] {
4987    // SAFETY: as dot_q4tp_row_1x8_avx512, two adjacent rows at once; the
4988    // caller guarantees r0 + 1 < rows and the ISA.
4989    unsafe {
4990        use core::arch::x86_64::*;
4991        let lomask = _mm256_set1_epi8(0x0F);
4992        let eight = _mm256_set1_epi8(8);
4993        let zero = _mm512_setzero_si512();
4994        let mut v0 = [_mm512_setzero_ps(); 8];
4995        let mut v1 = [_mm512_setzero_ps(); 8];
4996        let pairs = gpr / 2;
4997        let unpack = |r: usize, gi: usize| -> (__m512i, __mmask64) {
4998            let t = nib.as_ptr().add((r * gpr + gi) * Q4TP_NIB);
4999            let bb = _mm256_loadu_si256(t as *const __m256i);
5000            let lo = _mm256_and_si256(bb, lomask);
5001            let hi = _mm256_and_si256(_mm256_srli_epi16::<4>(bb), lomask);
5002            let ul = _mm256_sub_epi8(_mm256_unpacklo_epi8(lo, hi), eight);
5003            let uh = _mm256_sub_epi8(_mm256_unpackhi_epi8(lo, hi), eight);
5004            let cat = _mm512_inserti64x4::<1>(_mm512_castsi256_si512(ul), uh);
5005            let w = _mm512_shuffle_i64x2::<0b11_01_10_00>(cat, cat);
5006            (_mm512_abs_epi8(w), _mm512_movepi8_mask(w))
5007        };
5008        for gp in 0..pairs {
5009            let gi = gp * 2;
5010            let (wa0, neg0) = unpack(r0, gi);
5011            let (wa1, neg1) = unpack(r0 + 1, gi);
5012            let off = gi * GROUP_SIZE;
5013            let sv = |sc: &[f32]| {
5014                _mm512_insertf32x8::<1>(
5015                    _mm512_castps256_ps512(_mm256_set1_ps(*sc.get_unchecked(gi))),
5016                    _mm256_set1_ps(*sc.get_unchecked(gi + 1)),
5017                )
5018            };
5019            let s0 = sv(sc0);
5020            let s1 = sv(sc1);
5021            for k in 0..8 {
5022                let xv = _mm512_loadu_si512(xs[k].as_ptr().add(off) as *const __m512i);
5023                let d0 = _mm512_cvtepi32_ps(_mm512_dpbusd_epi32(
5024                    zero,
5025                    wa0,
5026                    _mm512_mask_sub_epi8(xv, neg0, zero, xv),
5027                ));
5028                let d1 = _mm512_cvtepi32_ps(_mm512_dpbusd_epi32(
5029                    zero,
5030                    wa1,
5031                    _mm512_mask_sub_epi8(xv, neg1, zero, xv),
5032                ));
5033                v0[k] = _mm512_fmadd_ps(d0, s0, v0[k]);
5034                v1[k] = _mm512_fmadd_ps(d1, s1, v1[k]);
5035            }
5036        }
5037        let mut acc = [[0f32; 8]; 2];
5038        for k in 0..8 {
5039            acc[0][k] = _mm512_reduce_add_ps(v0[k]);
5040            acc[1][k] = _mm512_reduce_add_ps(v1[k]);
5041        }
5042        if gpr % 2 == 1 {
5043            let off = (gpr - 1) * GROUP_SIZE;
5044            for j in off..off + GROUP_SIZE {
5045                let (w0, sa) = q4tp_outlier(nib, r0, gpr, j, sc0);
5046                let (w1, sb) = q4tp_outlier(nib, r0 + 1, gpr, j, sc1);
5047                for k in 0..8 {
5048                    let x = *xs[k].get_unchecked(j) as f32;
5049                    acc[0][k] += w0 * sa * x;
5050                    acc[1][k] += w1 * sb * x;
5051                }
5052            }
5053        }
5054        acc
5055    }
5056}
5057
5058/// The same, eight columns at a time. One unpack then feeds twice as many
5059/// activation streams, so a wide batch reads the weight tile half as
5060/// often; the price is eight accumulators live at once. Measured 9.0 ->
5061/// 8.3 ms at 9216x2304, b=296 on a 48-thread EPYC 9B45.
5062#[cfg(target_arch = "x86_64")]
5063#[target_feature(enable = "avx512f,avx512bw,avx512vnni")]
5064unsafe fn dot_q4tp_row_1x8_avx512(
5065    nib: &[u8],
5066    r: usize,
5067    gpr: usize,
5068    xs: [&[i8]; 8],
5069    scales: &[f32],
5070) -> [f32; 8] {
5071    // SAFETY: as dot_q4tp_row_1x4_avx2; caller guarantees the ISA.
5072    unsafe {
5073        use core::arch::x86_64::*;
5074        let lomask = _mm256_set1_epi8(0x0F);
5075        let eight = _mm256_set1_epi8(8);
5076        let zero = _mm512_setzero_si512();
5077        let (mut v0, mut v1, mut v2, mut v3) = (
5078            _mm512_setzero_ps(),
5079            _mm512_setzero_ps(),
5080            _mm512_setzero_ps(),
5081            _mm512_setzero_ps(),
5082        );
5083        let (mut v4, mut v5, mut v6, mut v7) = (
5084            _mm512_setzero_ps(),
5085            _mm512_setzero_ps(),
5086            _mm512_setzero_ps(),
5087            _mm512_setzero_ps(),
5088        );
5089        let pairs = gpr / 2;
5090        for gp in 0..pairs {
5091            let gi = gp * 2;
5092            let t = nib.as_ptr().add((r * gpr + gi) * Q4TP_NIB);
5093            let bb = _mm256_loadu_si256(t as *const __m256i);
5094            let lo = _mm256_and_si256(bb, lomask);
5095            let hi = _mm256_and_si256(_mm256_srli_epi16::<4>(bb), lomask);
5096            // `unpack` works per 128-bit lane, so the halves come out as
5097            // [A.lo, B.lo] and [A.hi, B.hi]; the shuffle reorders the four
5098            // 128-bit lanes into the weights' natural order, which is what
5099            // the straight activation load expects.
5100            let ul = _mm256_sub_epi8(_mm256_unpacklo_epi8(lo, hi), eight);
5101            let uh = _mm256_sub_epi8(_mm256_unpackhi_epi8(lo, hi), eight);
5102            let cat = _mm512_inserti64x4::<1>(_mm512_castsi256_si512(ul), uh);
5103            let w = _mm512_shuffle_i64x2::<0b11_01_10_00>(cat, cat);
5104            let wabs = _mm512_abs_epi8(w);
5105            let neg = _mm512_movepi8_mask(w);
5106            let off = gi * GROUP_SIZE;
5107            let sv = _mm512_insertf32x8::<1>(
5108                _mm512_castps256_ps512(_mm256_set1_ps(*scales.get_unchecked(gi))),
5109                _mm256_set1_ps(*scales.get_unchecked(gi + 1)),
5110            );
5111            let dot = |x: &[i8]| -> __m512 {
5112                let xv = _mm512_loadu_si512(x.as_ptr().add(off) as *const __m512i);
5113                let sx = _mm512_mask_sub_epi8(xv, neg, zero, xv);
5114                _mm512_cvtepi32_ps(_mm512_dpbusd_epi32(zero, wabs, sx))
5115            };
5116            v0 = _mm512_fmadd_ps(dot(xs[0]), sv, v0);
5117            v1 = _mm512_fmadd_ps(dot(xs[1]), sv, v1);
5118            v2 = _mm512_fmadd_ps(dot(xs[2]), sv, v2);
5119            v3 = _mm512_fmadd_ps(dot(xs[3]), sv, v3);
5120            v4 = _mm512_fmadd_ps(dot(xs[4]), sv, v4);
5121            v5 = _mm512_fmadd_ps(dot(xs[5]), sv, v5);
5122            v6 = _mm512_fmadd_ps(dot(xs[6]), sv, v6);
5123            v7 = _mm512_fmadd_ps(dot(xs[7]), sv, v7);
5124        }
5125        let mut acc = [
5126            _mm512_reduce_add_ps(v0),
5127            _mm512_reduce_add_ps(v1),
5128            _mm512_reduce_add_ps(v2),
5129            _mm512_reduce_add_ps(v3),
5130            _mm512_reduce_add_ps(v4),
5131            _mm512_reduce_add_ps(v5),
5132            _mm512_reduce_add_ps(v6),
5133            _mm512_reduce_add_ps(v7),
5134        ];
5135        // An odd group count leaves one group over; the narrow kernel
5136        // finishes it rather than the tail being a special case here.
5137        if gpr % 2 == 1 {
5138            let off = (gpr - 1) * GROUP_SIZE;
5139            for j in off..off + GROUP_SIZE {
5140                let (w, s) = q4tp_outlier(nib, r, gpr, j, scales);
5141                let ws = w * s;
5142                for k in 0..8 {
5143                    acc[k] += ws * *xs[k].get_unchecked(j) as f32;
5144                }
5145            }
5146        }
5147        acc
5148    }
5149}
5150
5151/// The same four columns, 512 bits wide. Two groups (64 weights) ride one
5152/// unpack and one `vpdpbusd`, where AVX2 needs two unpacks and four
5153/// `maddubs`/`madd` pairs — about 2.3x fewer instructions for the same
5154/// arithmetic. The two groups carry different scales, so the fma takes a
5155/// vector whose halves hold each group's scale rather than a broadcast.
5156///
5157/// There is no 512-bit `vpsignb`, so the activation's sign is applied by
5158/// negating under a mask taken from the weight's sign bits. That mask is
5159/// per-tile, so it is hoisted out of the column loop and the per-column
5160/// cost stays exactly one instruction, as with `sign_epi8`. Weights of
5161/// zero are not zeroed by the mask trick and do not need to be: their
5162/// magnitude is zero, so the product is.
5163#[cfg(target_arch = "x86_64")]
5164#[target_feature(enable = "avx512f,avx512bw,avx512vnni")]
5165unsafe fn dot_q4tp_row_1x4_avx512(
5166    nib: &[u8],
5167    r: usize,
5168    gpr: usize,
5169    xs: [&[i8]; 4],
5170    scales: &[f32],
5171) -> [f32; 4] {
5172    // SAFETY: as dot_q4tp_row_1x4_avx2; caller guarantees the ISA.
5173    unsafe {
5174        use core::arch::x86_64::*;
5175        let lomask = _mm256_set1_epi8(0x0F);
5176        let eight = _mm256_set1_epi8(8);
5177        let zero = _mm512_setzero_si512();
5178        let (mut v0, mut v1, mut v2, mut v3) = (
5179            _mm512_setzero_ps(),
5180            _mm512_setzero_ps(),
5181            _mm512_setzero_ps(),
5182            _mm512_setzero_ps(),
5183        );
5184        let pairs = gpr / 2;
5185        for gp in 0..pairs {
5186            let gi = gp * 2;
5187            let t = nib.as_ptr().add((r * gpr + gi) * Q4TP_NIB);
5188            let bb = _mm256_loadu_si256(t as *const __m256i);
5189            let lo = _mm256_and_si256(bb, lomask);
5190            let hi = _mm256_and_si256(_mm256_srli_epi16::<4>(bb), lomask);
5191            // `unpack` works per 128-bit lane, so the halves come out as
5192            // [A.lo, B.lo] and [A.hi, B.hi]; the shuffle reorders the four
5193            // 128-bit lanes into the weights' natural order, which is what
5194            // the straight activation load expects.
5195            let ul = _mm256_sub_epi8(_mm256_unpacklo_epi8(lo, hi), eight);
5196            let uh = _mm256_sub_epi8(_mm256_unpackhi_epi8(lo, hi), eight);
5197            let cat = _mm512_inserti64x4::<1>(_mm512_castsi256_si512(ul), uh);
5198            let w = _mm512_shuffle_i64x2::<0b11_01_10_00>(cat, cat);
5199            let wabs = _mm512_abs_epi8(w);
5200            let neg = _mm512_movepi8_mask(w);
5201            let off = gi * GROUP_SIZE;
5202            let sv = _mm512_insertf32x8::<1>(
5203                _mm512_castps256_ps512(_mm256_set1_ps(*scales.get_unchecked(gi))),
5204                _mm256_set1_ps(*scales.get_unchecked(gi + 1)),
5205            );
5206            let dot = |x: &[i8]| -> __m512 {
5207                let xv = _mm512_loadu_si512(x.as_ptr().add(off) as *const __m512i);
5208                let sx = _mm512_mask_sub_epi8(xv, neg, zero, xv);
5209                _mm512_cvtepi32_ps(_mm512_dpbusd_epi32(zero, wabs, sx))
5210            };
5211            v0 = _mm512_fmadd_ps(dot(xs[0]), sv, v0);
5212            v1 = _mm512_fmadd_ps(dot(xs[1]), sv, v1);
5213            v2 = _mm512_fmadd_ps(dot(xs[2]), sv, v2);
5214            v3 = _mm512_fmadd_ps(dot(xs[3]), sv, v3);
5215        }
5216        let mut acc = [
5217            _mm512_reduce_add_ps(v0),
5218            _mm512_reduce_add_ps(v1),
5219            _mm512_reduce_add_ps(v2),
5220            _mm512_reduce_add_ps(v3),
5221        ];
5222        // An odd group count leaves one group over; the narrow kernel
5223        // finishes it rather than the tail being a special case here.
5224        if gpr % 2 == 1 {
5225            let off = (gpr - 1) * GROUP_SIZE;
5226            for j in off..off + GROUP_SIZE {
5227                let (w, s) = q4tp_outlier(nib, r, gpr, j, scales);
5228                let ws = w * s;
5229                for k in 0..4 {
5230                    acc[k] += ws * *xs[k].get_unchecked(j) as f32;
5231                }
5232            }
5233        }
5234        acc
5235    }
5236}
5237
5238/// Four batch columns against one q4tp row: the tile is unpacked ONCE and
5239/// spent on four activation streams, which is where a prefill batch stops
5240/// being weight-bandwidth-bound. Twin of `dot_q4t_row_1x4_sdot`.
5241#[cfg(target_arch = "aarch64")]
5242#[target_feature(enable = "neon,dotprod")]
5243unsafe fn dot_q4tp_row_1x4_sdot(
5244    nib: &[u8],
5245    r: usize,
5246    gpr: usize,
5247    xs: [&[i8]; 4],
5248    scales: &[f32],
5249) -> [f32; 4] {
5250    // SAFETY: see dot_q4tp_row_sdot; every xs[k] is gpr·GROUP_SIZE long.
5251    unsafe {
5252        use core::arch::aarch64::*;
5253        use core::arch::asm;
5254        let lomask = vdupq_n_u8(0x0F);
5255        let eight = vdupq_n_s8(8);
5256        // Named accumulators, NOT an array indexed by a loop variable: the
5257        // latter does not stay in registers (the same defect cost 2x in the
5258        // AVX2 q4t kernel and again in WGSL).
5259        //
5260        // They are VECTORS, and the horizontal add happens once at the end
5261        // instead of once per group per column. `vaddvq` is a cross-lane
5262        // reduction — with 72 groups and four columns the old shape paid
5263        // 288 of them per row, each one a dependency stall the pipeline
5264        // cannot hide, to save four float adds. The group's scale now
5265        // rides an fma into the lane accumulators, so the arithmetic per
5266        // group is one convert and one fma. Summation order changes (the
5267        // lanes carry independent partial sums), which is the same
5268        // round-off class the SDOT path already lives in — the strict
5269        // kernel (`CMF_SDOT=0`, what `cortiq ppl` runs) is unchanged and
5270        // stays the reference.
5271        let (mut v0, mut v1, mut v2, mut v3) = (
5272            vdupq_n_f32(0.0),
5273            vdupq_n_f32(0.0),
5274            vdupq_n_f32(0.0),
5275            vdupq_n_f32(0.0),
5276        );
5277        for gi in 0..gpr {
5278            let t = nib.as_ptr().add((r * gpr + gi) * Q4TP_NIB);
5279            let s = *scales.get_unchecked(gi);
5280            let bb = vld1q_u8(t);
5281            let lo = vandq_u8(bb, lomask);
5282            let hi = vshrq_n_u8::<4>(bb);
5283            let e0 = vsubq_s8(vreinterpretq_s8_u8(vzip1q_u8(lo, hi)), eight);
5284            let e1 = vsubq_s8(vreinterpretq_s8_u8(vzip2q_u8(lo, hi)), eight);
5285            let off = gi * GROUP_SIZE;
5286            let dot4 = |x: &[i8]| -> int32x4_t {
5287                let x0 = vld1q_s8(x.as_ptr().add(off));
5288                let x1 = vld1q_s8(x.as_ptr().add(off + 16));
5289                let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
5290                asm!(
5291                    "sdot {a0:v}.4s, {e0:v}.16b, {x0:v}.16b",
5292                    "sdot {a1:v}.4s, {e1:v}.16b, {x1:v}.16b",
5293                    a0 = inout(vreg) a0, a1 = inout(vreg) a1,
5294                    e0 = in(vreg) e0, x0 = in(vreg) x0, e1 = in(vreg) e1, x1 = in(vreg) x1,
5295                    options(pure, nomem, nostack),
5296                );
5297                vaddq_s32(a0, a1)
5298            };
5299            v0 = vfmaq_n_f32(v0, vcvtq_f32_s32(dot4(xs[0])), s);
5300            v1 = vfmaq_n_f32(v1, vcvtq_f32_s32(dot4(xs[1])), s);
5301            v2 = vfmaq_n_f32(v2, vcvtq_f32_s32(dot4(xs[2])), s);
5302            v3 = vfmaq_n_f32(v3, vcvtq_f32_s32(dot4(xs[3])), s);
5303        }
5304        [
5305            vaddvq_f32(v0),
5306            vaddvq_f32(v1),
5307            vaddvq_f32(v2),
5308            vaddvq_f32(v3),
5309        ]
5310    }
5311}
5312
5313/// Fused q4tp matmat — the same three arms `q4t_matmat` has. Shipping only
5314/// the scalar one made Nanbeige-3B decode at 1.2 tok/s against q4t's 5.9:
5315/// the format was fine, the missing arms were the whole regression.
5316fn q4tp_matmat(
5317    bytes: &[u8],
5318    xs_all: &[f32],
5319    b: usize,
5320    rows: usize,
5321    cols: usize,
5322    out: &mut [f32],
5323    pool: Option<&Pool>,
5324) {
5325    debug_assert_eq!(out.len(), b * rows);
5326    let gpr = cols / GROUP_SIZE;
5327    let v = Q4tpView::new(bytes, rows, cols);
5328
5329    // Wide batches ride the AMX through a dequant-tile sgemm, as in q4t.
5330    #[cfg(target_os = "macos")]
5331    if b >= 8 && rows * cols >= 500_000 && accel_gemm_enabled() {
5332        dequant_matmat_accel(
5333            &|r, dst| {
5334                let mut sc = [0f32; 32];
5335                let mut scv;
5336                let s: &[f32] = if gpr <= 32 {
5337                    v.scales_into(r, gpr, &mut sc);
5338                    &sc[..gpr]
5339                } else {
5340                    scv = vec![0f32; gpr];
5341                    v.scales_into(r, gpr, &mut scv);
5342                    &scv
5343                };
5344                for gi in 0..gpr {
5345                    let tile = &v.nib[(r * gpr + gi) * Q4TP_NIB..(r * gpr + gi + 1) * Q4TP_NIB];
5346                    for (k, &bb) in tile.iter().enumerate() {
5347                        dst[gi * GROUP_SIZE + k * 2] = ((bb & 0x0F) as f32 - 8.0) * s[gi];
5348                        dst[gi * GROUP_SIZE + k * 2 + 1] =
5349                            (((bb >> 4) & 0x0F) as f32 - 8.0) * s[gi];
5350                    }
5351                }
5352            },
5353            xs_all,
5354            b,
5355            rows,
5356            cols,
5357            out,
5358            pool,
5359        );
5360        return;
5361    }
5362
5363    let out_addr = SendMut(out.as_mut_ptr());
5364    if a8w8_enabled() {
5365        let acts: Vec<SplitAct> = (0..b)
5366            .map(|bi| split_act(&xs_all[bi * cols..(bi + 1) * cols]))
5367            .collect();
5368        let acts = &acts;
5369        #[cfg(target_arch = "aarch64")]
5370        let blocked_ok = sdot_enabled() && blocked_enabled();
5371        // x86 gets the same blocking: one tile unpack spent on four
5372        // columns. Without it every column re-decoded the row, which is
5373        // why a 48-core EPYC measured a sixth of an M4's per-core rate.
5374        // The gate is `avx2_enabled`, as in q4t — `sdot_enabled` answers
5375        // for ARM's dotprod and is hard-wired false everywhere else, so
5376        // asking it here left the whole blocked path unreachable on x86.
5377        #[cfg(target_arch = "x86_64")]
5378        let blocked_ok = q4tp_blocked_x86();
5379        #[cfg(not(any(target_arch = "aarch64", target_arch = "x86_64")))]
5380        let blocked_ok = false;
5381        // Columns are swept in panels that fit L2. Without this a
5382        // row-pair walks every activation in the batch — 4.8 MB at
5383        // 512x512 — and does it again for the next pair, so the whole
5384        // batch streams out of the shared cache once per row. Measured
5385        // 800 GB/s of it, flat across batch sizes, which is the signature
5386        // of a loop bound by traffic rather than by arithmetic. A panel of
5387        // 256 columns is 590 KB beside 221 KB of this worker's weights:
5388        // both stay resident and the batch crosses L3 once instead of
5389        // once per row.
5390        let panel_cols: usize = std::env::var("CMF_Q4TP_PANEL")
5391            .ok()
5392            .and_then(|v| v.parse().ok())
5393            .filter(|v| *v > 0)
5394            .unwrap_or(256);
5395        let run = |start: usize, end: usize| {
5396            for abase in (0..acts.len()).step_by(panel_cols) {
5397                let alen = (acts.len() - abase).min(panel_cols);
5398                let mut sc = vec![0f32; gpr];
5399                #[cfg(target_arch = "x86_64")]
5400                let mut r_lo = start;
5401                #[cfg(target_arch = "x86_64")]
5402                if blocked_ok && alen >= 8 {
5403                    let mut sc1 = vec![0f32; gpr];
5404                    while r_lo + 2 <= end {
5405                        v.scales_into(r_lo, gpr, &mut sc);
5406                        v.scales_into(r_lo + 1, gpr, &mut sc1);
5407                        let mut bi = 0usize;
5408                        while bi + 8 <= alen {
5409                            let xs = [
5410                                acts[abase + bi].xq.as_slice(),
5411                                acts[abase + bi + 1].xq.as_slice(),
5412                                acts[abase + bi + 2].xq.as_slice(),
5413                                acts[abase + bi + 3].xq.as_slice(),
5414                                acts[abase + bi + 4].xq.as_slice(),
5415                                acts[abase + bi + 5].xq.as_slice(),
5416                                acts[abase + bi + 6].xq.as_slice(),
5417                                acts[abase + bi + 7].xq.as_slice(),
5418                            ];
5419                            let d = unsafe { dot_q4tp_2x8_avx512(v.nib, r_lo, gpr, xs, &sc, &sc1) };
5420                            for (row, dr, scr) in [(r_lo, &d[0], &sc), (r_lo + 1, &d[1], &sc1)] {
5421                                for k in 0..8 {
5422                                    let act = &acts[abase + bi + k];
5423                                    let mut acc = dr[k] * act.sx;
5424                                    for &(j, xv) in &act.outliers {
5425                                        let (w, s) = q4tp_outlier(v.nib, row, gpr, j, scr);
5426                                        acc += w * s * xv;
5427                                    }
5428                                    // SAFETY: disjoint (bi, r) cells per worker.
5429                                    unsafe { *out_addr.at((abase + bi + k) * rows + row) = acc };
5430                                }
5431                            }
5432                            bi += 8;
5433                        }
5434                        // Columns past the last group of eight, both rows —
5435                        // the same single-row kernel the tail below uses.
5436                        for row in [r_lo, r_lo + 1] {
5437                            let scr: &[f32] = if row == r_lo { &sc } else { &sc1 };
5438                            for b2 in bi..alen {
5439                                let act = &acts[abase + b2];
5440                                let xs4 = [
5441                                    act.xq.as_slice(),
5442                                    act.xq.as_slice(),
5443                                    act.xq.as_slice(),
5444                                    act.xq.as_slice(),
5445                                ];
5446                                let d =
5447                                    unsafe { dot_q4tp_row_1x4_avx512(v.nib, row, gpr, xs4, scr) };
5448                                let mut acc = d[0] * act.sx;
5449                                for &(j, xv) in &act.outliers {
5450                                    let (w, s) = q4tp_outlier(v.nib, row, gpr, j, scr);
5451                                    acc += w * s * xv;
5452                                }
5453                                // SAFETY: disjoint (bi, r) cells per worker.
5454                                unsafe { *out_addr.at((abase + b2) * rows + row) = acc };
5455                            }
5456                        }
5457                        r_lo += 2;
5458                    }
5459                }
5460                #[cfg(target_arch = "x86_64")]
5461                let row_start = r_lo;
5462                #[cfg(not(target_arch = "x86_64"))]
5463                let row_start = start;
5464                for r in row_start..end {
5465                    v.scales_into(r, gpr, &mut sc);
5466                    let mut bi = 0usize;
5467                    #[cfg(target_arch = "x86_64")]
5468                    if blocked_ok {
5469                        while bi + 8 <= alen {
5470                            let xs = [
5471                                acts[abase + bi].xq.as_slice(),
5472                                acts[abase + bi + 1].xq.as_slice(),
5473                                acts[abase + bi + 2].xq.as_slice(),
5474                                acts[abase + bi + 3].xq.as_slice(),
5475                                acts[abase + bi + 4].xq.as_slice(),
5476                                acts[abase + bi + 5].xq.as_slice(),
5477                                acts[abase + bi + 6].xq.as_slice(),
5478                                acts[abase + bi + 7].xq.as_slice(),
5479                            ];
5480                            let d = unsafe { dot_q4tp_row_1x8_avx512(v.nib, r, gpr, xs, &sc) };
5481                            for k in 0..8 {
5482                                let act = &acts[abase + bi + k];
5483                                let mut acc = d[k] * act.sx;
5484                                for &(j, xv) in &act.outliers {
5485                                    let (w, s) = q4tp_outlier(v.nib, r, gpr, j, &sc);
5486                                    acc += w * s * xv;
5487                                }
5488                                // SAFETY: disjoint (bi, r) cells per worker.
5489                                unsafe { *out_addr.at((abase + bi + k) * rows + r) = acc };
5490                            }
5491                            bi += 8;
5492                        }
5493                        while bi + 4 <= alen {
5494                            let xs = [
5495                                acts[abase + bi].xq.as_slice(),
5496                                acts[abase + bi + 1].xq.as_slice(),
5497                                acts[abase + bi + 2].xq.as_slice(),
5498                                acts[abase + bi + 3].xq.as_slice(),
5499                            ];
5500                            let d = unsafe { dot_q4tp_row_1x4_avx512(v.nib, r, gpr, xs, &sc) };
5501                            for k in 0..4 {
5502                                let act = &acts[abase + bi + k];
5503                                let mut acc = d[k] * act.sx;
5504                                for &(j, xv) in &act.outliers {
5505                                    let (w, s) = q4tp_outlier(v.nib, r, gpr, j, &sc);
5506                                    acc += w * s * xv;
5507                                }
5508                                // SAFETY: disjoint (bi, r) cells per worker.
5509                                unsafe { *out_addr.at((abase + bi + k) * rows + r) = acc };
5510                            }
5511                            bi += 4;
5512                        }
5513                    }
5514                    #[cfg(target_arch = "aarch64")]
5515                    if blocked_ok {
5516                        while bi + 4 <= alen {
5517                            let xs = [
5518                                acts[abase + bi].xq.as_slice(),
5519                                acts[abase + bi + 1].xq.as_slice(),
5520                                acts[abase + bi + 2].xq.as_slice(),
5521                                acts[abase + bi + 3].xq.as_slice(),
5522                            ];
5523                            let d = unsafe {
5524                                if q4tp_v1() {
5525                                    dot_q4tp_row_1x4_sdot_v1(v.nib, r, gpr, xs, &sc)
5526                                } else {
5527                                    dot_q4tp_row_1x4_sdot(v.nib, r, gpr, xs, &sc)
5528                                }
5529                            };
5530                            for k in 0..4 {
5531                                let act = &acts[abase + bi + k];
5532                                let mut acc = d[k] * act.sx;
5533                                for &(j, xv) in &act.outliers {
5534                                    let (w, s) = q4tp_outlier(v.nib, r, gpr, j, &sc);
5535                                    acc += w * s * xv;
5536                                }
5537                                // SAFETY: disjoint (bi, r) cells per worker.
5538                                unsafe { *out_addr.at((abase + bi + k) * rows + r) = acc };
5539                            }
5540                            bi += 4;
5541                        }
5542                    }
5543                    let _ = blocked_ok;
5544                    while bi < alen {
5545                        let act = &acts[abase + bi];
5546                        let mut acc = dot_q4tp_row_i8(v.nib, r, gpr, &act.xq, &sc) * act.sx;
5547                        for &(j, xv) in &act.outliers {
5548                            let (w, s) = q4tp_outlier(v.nib, r, gpr, j, &sc);
5549                            acc += w * s * xv;
5550                        }
5551                        // SAFETY: disjoint (bi, r) cells per worker range.
5552                        unsafe { *out_addr.at((abase + bi) * rows + r) = acc };
5553                        bi += 1;
5554                    }
5555                }
5556            }
5557        };
5558        dispatch_rows(pool, rows, &run);
5559        return;
5560    }
5561
5562    let run = |start: usize, end: usize| {
5563        let mut sc = vec![0f32; gpr];
5564        for r in start..end {
5565            v.scales_into(r, gpr, &mut sc);
5566            for bi in 0..b {
5567                let x = &xs_all[bi * cols..(bi + 1) * cols];
5568                // SAFETY: disjoint (bi, r) cells per worker range.
5569                unsafe { *out_addr.at(bi * rows + r) = q4tp_row_exact(v.nib, r, gpr, x, &sc) };
5570            }
5571        }
5572    };
5573    dispatch_rows(pool, rows, &run);
5574}
5575
5576/// Fused q4_tiled matvec (dispatch mirrors `q4matvec`).
5577fn q4t_matvec(
5578    bytes: &[u8],
5579    x: &[f32],
5580    rows: usize,
5581    cols: usize,
5582    out: &mut [f32],
5583    pool: Option<&Pool>,
5584) {
5585    debug_assert_eq!(out.len(), rows);
5586    let gpr = cols / GROUP_SIZE;
5587    let out_addr = SendMut(out.as_mut_ptr());
5588    if a8w8_enabled() {
5589        let act = split_act(x);
5590        let run = move |start: usize, end: usize| {
5591            for r in start..end {
5592                let mut acc = dot_q4t_row_i8(bytes, r, gpr, &act.xq) * act.sx;
5593                for &(j, xv) in &act.outliers {
5594                    let (w, s) = q4t_outlier(bytes, r, gpr, j);
5595                    acc += w * s * xv;
5596                }
5597                // SAFETY: disjoint row ranges per worker.
5598                unsafe { *out_addr.at(r) = acc };
5599            }
5600        };
5601        dispatch_rows(pool, rows, &run);
5602        return;
5603    }
5604    let run = move |start: usize, end: usize| {
5605        for r in start..end {
5606            // SAFETY: disjoint row ranges per worker.
5607            unsafe { *out_addr.at(r) = q4t_row_exact(bytes, r, gpr, x) };
5608        }
5609    };
5610    dispatch_rows(pool, rows, &run);
5611}
5612
5613/// Fused two-input q4_tiled matvec (weights read once per pair).
5614#[allow(clippy::too_many_arguments)]
5615fn q4t_matvec2(
5616    bytes: &[u8],
5617    x1: &[f32],
5618    x2: &[f32],
5619    rows: usize,
5620    cols: usize,
5621    o1: &mut [f32],
5622    o2: &mut [f32],
5623    pool: Option<&Pool>,
5624) {
5625    let gpr = cols / GROUP_SIZE;
5626    let p1 = SendMut(o1.as_mut_ptr());
5627    let p2 = SendMut(o2.as_mut_ptr());
5628    if a8w8_enabled() {
5629        let a1 = split_act(x1);
5630        let a2 = split_act(x2);
5631        let run = move |start: usize, end: usize| {
5632            for r in start..end {
5633                let mut v1 = dot_q4t_row_i8(bytes, r, gpr, &a1.xq) * a1.sx;
5634                let mut v2 = dot_q4t_row_i8(bytes, r, gpr, &a2.xq) * a2.sx;
5635                for &(j, xv) in &a1.outliers {
5636                    let (w, s) = q4t_outlier(bytes, r, gpr, j);
5637                    v1 += w * s * xv;
5638                }
5639                for &(j, xv) in &a2.outliers {
5640                    let (w, s) = q4t_outlier(bytes, r, gpr, j);
5641                    v2 += w * s * xv;
5642                }
5643                // SAFETY: disjoint row ranges per worker.
5644                unsafe {
5645                    *p1.at(r) = v1;
5646                    *p2.at(r) = v2;
5647                }
5648            }
5649        };
5650        dispatch_rows(pool, rows, &run);
5651        return;
5652    }
5653    let run = move |start: usize, end: usize| {
5654        for r in start..end {
5655            // SAFETY: disjoint row ranges per worker.
5656            unsafe {
5657                *p1.at(r) = q4t_row_exact(bytes, r, gpr, x1);
5658                *p2.at(r) = q4t_row_exact(bytes, r, gpr, x2);
5659            }
5660        }
5661    };
5662    dispatch_rows(pool, rows, &run);
5663}
5664
5665/// Batched q4_tiled matmat: each row's tiles stream once per microbatch.
5666#[allow(clippy::too_many_arguments)]
5667/// Prefill GEMM through Accelerate for group-quantized codecs: a
5668/// caller-supplied row dequantizer fills f32 tiles (pool-parallel) and
5669/// each tile rides the AMX with one sgemm — the generic sibling of
5670/// `qmatmat_accel` (q8). Numerics are f32-GEMM (tolerance class);
5671/// decode (b=1) never takes this path.
5672#[cfg(target_os = "macos")]
5673fn dequant_matmat_accel(
5674    dequant_row: &(dyn Fn(usize, &mut [f32]) + Sync),
5675    xs_all: &[f32],
5676    b: usize,
5677    rows: usize,
5678    cols: usize,
5679    out: &mut [f32],
5680    pool: Option<&Pool>,
5681) {
5682    const TR: usize = 2048;
5683    thread_local! {
5684        static WTILE: std::cell::RefCell<Vec<f32>> = const { std::cell::RefCell::new(Vec::new()) };
5685    }
5686    WTILE.with(|wt| {
5687        let mut wtile = wt.borrow_mut();
5688        wtile.resize(TR * cols, 0.0);
5689        let mut r0 = 0usize;
5690        while r0 < rows {
5691            let tr = TR.min(rows - r0);
5692            let wt_addr = SendMut(wtile.as_mut_ptr());
5693            let run = |start: usize, end: usize| {
5694                for r in start..end {
5695                    // SAFETY: workers cover disjoint r ranges.
5696                    let dst = unsafe { std::slice::from_raw_parts_mut(wt_addr.at(r * cols), cols) };
5697                    dequant_row(r0 + r, dst);
5698                }
5699            };
5700            dispatch_rows(pool, tr, &run);
5701            unsafe {
5702                accel_blas::cblas_sgemm(
5703                    101, // RowMajor
5704                    111, // NoTrans A
5705                    112, // Trans B
5706                    b as i32,
5707                    tr as i32,
5708                    cols as i32,
5709                    1.0,
5710                    xs_all.as_ptr(),
5711                    cols as i32,
5712                    wtile.as_ptr(),
5713                    cols as i32,
5714                    0.0,
5715                    out.as_mut_ptr().add(r0),
5716                    rows as i32,
5717                );
5718            }
5719            r0 += tr;
5720        }
5721    });
5722}
5723
5724fn q4t_matmat(
5725    bytes: &[u8],
5726    xs_all: &[f32],
5727    b: usize,
5728    rows: usize,
5729    cols: usize,
5730    out: &mut [f32],
5731    pool: Option<&Pool>,
5732) {
5733    debug_assert_eq!(out.len(), b * rows);
5734    let gpr = cols / GROUP_SIZE;
5735    // Wide batches ride the AMX like q8's qmatmat: on Apple silicon
5736    // the dequant-tile sgemm is an order above the SDOT row loop for
5737    // prefill shapes (imagegen DiT forwards are exactly this).
5738    #[cfg(target_os = "macos")]
5739    if b >= 8 && rows * cols >= 500_000 && accel_gemm_enabled() {
5740        dequant_matmat_accel(
5741            &|r, dst| {
5742                for gi in 0..gpr {
5743                    let tile = &bytes[(r * gpr + gi) * Q4_TILE..(r * gpr + gi + 1) * Q4_TILE];
5744                    let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
5745                    for (k, &bb) in tile[2..].iter().enumerate() {
5746                        dst[gi * GROUP_SIZE + k * 2] = ((bb & 0x0F) as f32 - 8.0) * s;
5747                        dst[gi * GROUP_SIZE + k * 2 + 1] = (((bb >> 4) & 0x0F) as f32 - 8.0) * s;
5748                    }
5749                }
5750            },
5751            xs_all,
5752            b,
5753            rows,
5754            cols,
5755            out,
5756            pool,
5757        );
5758        return;
5759    }
5760    let out_addr = SendMut(out.as_mut_ptr());
5761    if a8w8_enabled() {
5762        let acts: Vec<SplitAct> = (0..b)
5763            .map(|bi| split_act(&xs_all[bi * cols..(bi + 1) * cols]))
5764            .collect();
5765        let acts = &acts;
5766        #[cfg(target_arch = "x86_64")]
5767        let blocked_ok = avx2_enabled() && blocked_enabled();
5768        #[cfg(target_arch = "aarch64")]
5769        let blocked_ok = sdot_enabled() && blocked_enabled();
5770        #[cfg(not(any(target_arch = "x86_64", target_arch = "aarch64")))]
5771        let blocked_ok = false;
5772        let run = move |start: usize, end: usize| {
5773            for r in start..end {
5774                let mut bi = 0usize;
5775                #[cfg(target_arch = "aarch64")]
5776                if blocked_ok {
5777                    while bi + 4 <= acts.len() {
5778                        let xs = [
5779                            acts[bi].xq.as_slice(),
5780                            acts[bi + 1].xq.as_slice(),
5781                            acts[bi + 2].xq.as_slice(),
5782                            acts[bi + 3].xq.as_slice(),
5783                        ];
5784                        let d = unsafe { dot_q4t_row_1x4_sdot(bytes, r, gpr, xs) };
5785                        for k in 0..4 {
5786                            let act = &acts[bi + k];
5787                            let mut acc = d[k] * act.sx;
5788                            for &(j, xv) in &act.outliers {
5789                                let (w, sc) = q4t_outlier(bytes, r, gpr, j);
5790                                acc += w * sc * xv;
5791                            }
5792                            // SAFETY: disjoint (bi, r) cells per worker.
5793                            unsafe { *out_addr.at((bi + k) * rows + r) = acc };
5794                        }
5795                        bi += 4;
5796                    }
5797                }
5798                #[cfg(target_arch = "x86_64")]
5799                if blocked_ok {
5800                    while bi + 4 <= acts.len() {
5801                        let xs = [
5802                            acts[bi].xq.as_slice(),
5803                            acts[bi + 1].xq.as_slice(),
5804                            acts[bi + 2].xq.as_slice(),
5805                            acts[bi + 3].xq.as_slice(),
5806                        ];
5807                        let d = unsafe {
5808                            if vnni_tiles_enabled() {
5809                                dot_q4t_row_1x4_vnni(bytes, r, gpr, xs)
5810                            } else {
5811                                dot_q4t_row_1x4_avx2(bytes, r, gpr, xs)
5812                            }
5813                        };
5814                        for k in 0..4 {
5815                            let act = &acts[bi + k];
5816                            let mut acc = d[k] * act.sx;
5817                            for &(j, xv) in &act.outliers {
5818                                let (w, sc) = q4t_outlier(bytes, r, gpr, j);
5819                                acc += w * sc * xv;
5820                            }
5821                            // SAFETY: disjoint (bi, r) cells per worker.
5822                            unsafe { *out_addr.at((bi + k) * rows + r) = acc };
5823                        }
5824                        bi += 4;
5825                    }
5826                }
5827                let _ = blocked_ok;
5828                while bi < acts.len() {
5829                    let act = &acts[bi];
5830                    let mut acc = dot_q4t_row_i8(bytes, r, gpr, &act.xq) * act.sx;
5831                    for &(j, xv) in &act.outliers {
5832                        let (w, s) = q4t_outlier(bytes, r, gpr, j);
5833                        acc += w * s * xv;
5834                    }
5835                    // SAFETY: disjoint (bi, r) cells per worker range.
5836                    unsafe { *out_addr.at(bi * rows + r) = acc };
5837                    bi += 1;
5838                }
5839            }
5840        };
5841        dispatch_rows(pool, rows, &run);
5842        return;
5843    }
5844    let run = move |start: usize, end: usize| {
5845        for r in start..end {
5846            for bi in 0..b {
5847                let x = &xs_all[bi * cols..(bi + 1) * cols];
5848                // SAFETY: disjoint (bi, r) cells per worker range.
5849                unsafe { *out_addr.at(bi * rows + r) = q4t_row_exact(bytes, r, gpr, x) };
5850            }
5851        }
5852    };
5853    dispatch_rows(pool, rows, &run);
5854}
5855
5856// ── q1 (dtype 12): binary weights, [f16 scale][4B sign bits] per
5857// 32-group tile. The kernel family mirrors q4_tiled: one sequential
5858// stream of 6-byte tiles, per-tile integer dot × scale, exact outlier
5859// correction (A8W8 contract), exact scalar path under CMF_SDOT=0. ──
5860
5861/// Per-32-group sums of the quantized activation — the ±1 identity's
5862/// shared half: `dot = −2·sdot(mask, x) − gsum[g]`, computed ONCE per
5863/// matvec and reused by every row.
5864fn q1_group_sums(xq: &[i8], gpr: usize) -> Vec<i32> {
5865    (0..gpr)
5866        .map(|gi| {
5867            xq[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE]
5868                .iter()
5869                .map(|&v| v as i32)
5870                .sum()
5871        })
5872        .collect()
5873}
5874
5875/// One q1 row via the A8W8 int8 path — mask-SDOT on ARM (no ±1
5876/// expansion at all), scalar bit loop elsewhere (AVX2 queued with the
5877/// x86 pass).
5878#[inline]
5879#[allow(unreachable_code)]
5880/// AVX2 q1 row via the same ±1 identity as the ARM sdot kernel: the
5881/// sign bits expand to a {0, −1} byte mask through shuffle+cmpeq, the
5882/// masked activation sums through maddubs(1, x&mask), and
5883/// `dot = −(2·masked_sum + Σx_group)` — bit-identical integer math.
5884#[cfg(target_arch = "x86_64")]
5885#[target_feature(enable = "avx2")]
5886unsafe fn dot_q1_row_avx2(bytes: &[u8], r: usize, gpr: usize, xq: &[i8], gsum: &[i32]) -> f32 {
5887    // SAFETY: callers uphold the 6B-tile and xq/gsum length contracts.
5888    unsafe {
5889        use core::arch::x86_64::*;
5890        // Byte j of the mask must replicate bits-byte j/8.
5891        let expand = _mm256_setr_epi8(
5892            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,
5893            3, 3, 3,
5894        );
5895        let bitsel = _mm256_setr_epi8(
5896            1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64,
5897            -128, 1, 2, 4, 8, 16, 32, 64, -128,
5898        );
5899        let ones8 = _mm256_set1_epi8(1);
5900        let ones16 = _mm256_set1_epi16(1);
5901        let mut acc = 0f32;
5902        for gi in 0..gpr {
5903            let t = bytes.as_ptr().add((r * gpr + gi) * Q1_TILE);
5904            let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
5905            let bits = u32::from_le_bytes([*t.add(2), *t.add(3), *t.add(4), *t.add(5)]);
5906            let bc = _mm256_shuffle_epi8(_mm256_set1_epi32(bits as i32), expand);
5907            let mask = _mm256_cmpeq_epi8(_mm256_and_si256(bc, bitsel), bitsel);
5908            let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
5909            let sel = _mm256_and_si256(x, mask);
5910            // Σ of selected i8 lanes: maddubs(1u8, sel_i8) pairs → madd.
5911            let p16 = _mm256_maddubs_epi16(ones8, sel);
5912            let d32 = _mm256_madd_epi16(p16, ones16);
5913            let hi128 = _mm256_extracti128_si256::<1>(d32);
5914            let s128 = _mm_add_epi32(_mm256_castsi256_si128(d32), hi128);
5915            let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
5916            let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
5917            let msum = _mm_cvtsi128_si32(s32);
5918            // The and-select keeps x UN-negated (unlike ARM's −1-mask
5919            // sdot): d = Σ_set − Σ_unset = 2·Σ_set − Σ_all.
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_avx2`: the masked-select sum goes through
5928/// one `vpdpbusd(1u8, sel)` (see `dpbusd_hsum` — bit-identical).
5929#[cfg(target_arch = "x86_64")]
5930#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
5931unsafe fn dot_q1_row_vnni(bytes: &[u8], r: usize, gpr: usize, xq: &[i8], gsum: &[i32]) -> f32 {
5932    // SAFETY: callers uphold the 6B-tile and xq/gsum length contracts.
5933    unsafe {
5934        use core::arch::x86_64::*;
5935        let expand = _mm256_setr_epi8(
5936            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,
5937            3, 3, 3,
5938        );
5939        let bitsel = _mm256_setr_epi8(
5940            1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64,
5941            -128, 1, 2, 4, 8, 16, 32, 64, -128,
5942        );
5943        let ones8 = _mm256_set1_epi8(1);
5944        let mut acc = 0f32;
5945        for gi in 0..gpr {
5946            let t = bytes.as_ptr().add((r * gpr + gi) * Q1_TILE);
5947            let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
5948            let bits = u32::from_le_bytes([*t.add(2), *t.add(3), *t.add(4), *t.add(5)]);
5949            let bc = _mm256_shuffle_epi8(_mm256_set1_epi32(bits as i32), expand);
5950            let mask = _mm256_cmpeq_epi8(_mm256_and_si256(bc, bitsel), bitsel);
5951            let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
5952            let msum = dpbusd_hsum(ones8, _mm256_and_si256(x, mask));
5953            let d = 2 * msum - gsum[gi];
5954            acc += d as f32 * s;
5955        }
5956        acc
5957    }
5958}
5959
5960/// VNNI twin of `dot_q1_row_1x4_avx2` (see `dpbusd_hsum`).
5961#[cfg(target_arch = "x86_64")]
5962#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
5963unsafe fn dot_q1_row_1x4_vnni(
5964    bytes: &[u8],
5965    r: usize,
5966    gpr: usize,
5967    xs: [&[i8]; 4],
5968    gsums: [&[i32]; 4],
5969) -> [f32; 4] {
5970    // SAFETY: callers uphold the 6B-tile and xq/gsum length contracts.
5971    unsafe {
5972        use core::arch::x86_64::*;
5973        let expand = _mm256_setr_epi8(
5974            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,
5975            3, 3, 3,
5976        );
5977        let bitsel = _mm256_setr_epi8(
5978            1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64,
5979            -128, 1, 2, 4, 8, 16, 32, 64, -128,
5980        );
5981        let ones8 = _mm256_set1_epi8(1);
5982        let mut acc = [0f32; 4];
5983        for gi in 0..gpr {
5984            let t = bytes.as_ptr().add((r * gpr + gi) * Q1_TILE);
5985            let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
5986            let bits = u32::from_le_bytes([*t.add(2), *t.add(3), *t.add(4), *t.add(5)]);
5987            let bc = _mm256_shuffle_epi8(_mm256_set1_epi32(bits as i32), expand);
5988            let mask = _mm256_cmpeq_epi8(_mm256_and_si256(bc, bitsel), bitsel);
5989            for (k, xq) in xs.iter().enumerate() {
5990                let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
5991                let msum = dpbusd_hsum(ones8, _mm256_and_si256(x, mask));
5992                let d = 2 * msum - gsums[k][gi];
5993                acc[k] += d as f32 * s;
5994            }
5995        }
5996        acc
5997    }
5998}
5999
6000/// The blocked 1×4 flavor: the expanded bit mask serves four activation
6001/// streams per group (mask build once, four select+reduce chains).
6002#[cfg(target_arch = "x86_64")]
6003#[target_feature(enable = "avx2")]
6004unsafe fn dot_q1_row_1x4_avx2(
6005    bytes: &[u8],
6006    r: usize,
6007    gpr: usize,
6008    xs: [&[i8]; 4],
6009    gsums: [&[i32]; 4],
6010) -> [f32; 4] {
6011    // SAFETY: callers uphold the 6B-tile and xq/gsum length contracts.
6012    unsafe {
6013        use core::arch::x86_64::*;
6014        let expand = _mm256_setr_epi8(
6015            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,
6016            3, 3, 3,
6017        );
6018        let bitsel = _mm256_setr_epi8(
6019            1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64,
6020            -128, 1, 2, 4, 8, 16, 32, 64, -128,
6021        );
6022        let ones8 = _mm256_set1_epi8(1);
6023        let ones16 = _mm256_set1_epi16(1);
6024        let mut acc = [0f32; 4];
6025        for gi in 0..gpr {
6026            let t = bytes.as_ptr().add((r * gpr + gi) * Q1_TILE);
6027            let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
6028            let bits = u32::from_le_bytes([*t.add(2), *t.add(3), *t.add(4), *t.add(5)]);
6029            let bc = _mm256_shuffle_epi8(_mm256_set1_epi32(bits as i32), expand);
6030            let mask = _mm256_cmpeq_epi8(_mm256_and_si256(bc, bitsel), bitsel);
6031            for (k, xq) in xs.iter().enumerate() {
6032                let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
6033                let sel = _mm256_and_si256(x, mask);
6034                let p16 = _mm256_maddubs_epi16(ones8, sel);
6035                let d32 = _mm256_madd_epi16(p16, ones16);
6036                let hi128 = _mm256_extracti128_si256::<1>(d32);
6037                let s128 = _mm_add_epi32(_mm256_castsi256_si128(d32), hi128);
6038                let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
6039                let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
6040                let msum = _mm_cvtsi128_si32(s32);
6041                let d = 2 * msum - gsums[k][gi];
6042                acc[k] += d as f32 * s;
6043            }
6044        }
6045        acc
6046    }
6047}
6048
6049#[allow(unreachable_code)]
6050fn dot_q1_row_i8(bytes: &[u8], r: usize, gpr: usize, xq: &[i8], gsum: &[i32]) -> f32 {
6051    #[cfg(target_arch = "aarch64")]
6052    unsafe {
6053        return dot_q1_row_sdot(bytes, r, gpr, xq, gsum);
6054    }
6055    #[cfg(target_arch = "x86_64")]
6056    if avx2_enabled() {
6057        unsafe {
6058            if vnni_tiles_enabled() {
6059                return dot_q1_row_vnni(bytes, r, gpr, xq, gsum);
6060            }
6061            return dot_q1_row_avx2(bytes, r, gpr, xq, gsum);
6062        }
6063    }
6064    let _ = gsum;
6065    let mut acc = 0f32;
6066    for gi in 0..gpr {
6067        let tile = &bytes[(r * gpr + gi) * Q1_TILE..(r * gpr + gi + 1) * Q1_TILE];
6068        let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
6069        let mut d = 0i32;
6070        for (j, &b) in tile[2..].iter().enumerate() {
6071            for k in 0..8 {
6072                let w = ((b >> k) & 1) as i32 * 2 - 1;
6073                d += w * xq[gi * GROUP_SIZE + j * 8 + k] as i32;
6074            }
6075        }
6076        acc += d as f32 * s;
6077    }
6078    acc
6079}
6080
6081/// SDOT q1 row via the ±1 identity: the vtst mask (0xFF where the bit
6082/// is set, i.e. −1 as i8) feeds `sdot` DIRECTLY — no expansion to ±1
6083/// lanes at all — and `dot = −(2·sdot(mask, x) + Σx_group)`, with the
6084/// per-group activation sums shared across every row of the matvec.
6085/// Four tiles (128 weights) per iteration: integer dots reduce through
6086/// a vpaddq tree into ONE i32x4 that meets its four scales in a single
6087/// fused f32 multiply-add. Integer math throughout — bit-identical to
6088/// the scalar ±1 reference.
6089#[cfg(target_arch = "aarch64")]
6090#[target_feature(enable = "neon,dotprod")]
6091unsafe fn dot_q1_row_sdot(bytes: &[u8], r: usize, gpr: usize, xq: &[i8], gsum: &[i32]) -> f32 {
6092    // SAFETY: callers uphold slice-length contracts (6B tile per group,
6093    // xq.len() == gpr·GROUP_SIZE, gsum.len() == gpr).
6094    unsafe {
6095        use core::arch::aarch64::*;
6096        use core::arch::asm;
6097        const MASKS: [u8; 16] = [1, 2, 4, 8, 16, 32, 64, 128, 1, 2, 4, 8, 16, 32, 64, 128];
6098        let m = vld1q_u8(MASKS.as_ptr());
6099        // One tile's −Σ_set(x) as an UNREDUCED i32x4 (two mask-sdots).
6100        macro_rules! tile_dot {
6101            ($t:expr, $x:expr) => {{
6102                let v0 = vcombine_u8(vdup_n_u8(*$t.add(2)), vdup_n_u8(*$t.add(3)));
6103                let v1 = vcombine_u8(vdup_n_u8(*$t.add(4)), vdup_n_u8(*$t.add(5)));
6104                let w0 = vreinterpretq_s8_u8(vtstq_u8(v0, m));
6105                let w1 = vreinterpretq_s8_u8(vtstq_u8(v1, m));
6106                let x0 = vld1q_s8($x);
6107                let x1 = vld1q_s8($x.add(16));
6108                let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
6109                asm!(
6110                    "sdot {a0:v}.4s, {w0:v}.16b, {x0:v}.16b",
6111                    "sdot {a1:v}.4s, {w1:v}.16b, {x1:v}.16b",
6112                    a0 = inout(vreg) a0, a1 = inout(vreg) a1,
6113                    w0 = in(vreg) w0, x0 = in(vreg) x0, w1 = in(vreg) w1, x1 = in(vreg) x1,
6114                    options(pure, nomem, nostack),
6115                );
6116                vaddq_s32(a0, a1)
6117            }};
6118        }
6119        // TBL unpack over PAIR loads: one vld1q covers two 6B tiles
6120        // ([s s b b b b][s s b b b b] + 4B slack), TBL replicates each
6121        // bit-byte across 8 lanes for vtst, and the four scales gather
6122        // through tbl2 into one fcvtl — the 16 ld1r broadcast loads and
6123        // 4 branchy software f16 conversions per 128 weights (the
6124        // measured load-port wall of this kernel) become 2 vector
6125        // loads + 9 table lookups. Integer math order is unchanged —
6126        // bit-identical results (FCVTL is exact on every f16).
6127        const IW00: [u8; 16] = [2, 2, 2, 2, 2, 2, 2, 2, 3, 3, 3, 3, 3, 3, 3, 3];
6128        const IW01: [u8; 16] = [4, 4, 4, 4, 4, 4, 4, 4, 5, 5, 5, 5, 5, 5, 5, 5];
6129        const IW10: [u8; 16] = [8, 8, 8, 8, 8, 8, 8, 8, 9, 9, 9, 9, 9, 9, 9, 9];
6130        const IW11: [u8; 16] = [
6131            10, 10, 10, 10, 10, 10, 10, 10, 11, 11, 11, 11, 11, 11, 11, 11,
6132        ];
6133        const ISC: [u8; 8] = [0, 1, 6, 7, 16, 17, 22, 23];
6134        let (iw00, iw01) = (vld1q_u8(IW00.as_ptr()), vld1q_u8(IW01.as_ptr()));
6135        let (iw10, iw11) = (vld1q_u8(IW10.as_ptr()), vld1q_u8(IW11.as_ptr()));
6136        let isc = vld1_u8(ISC.as_ptr());
6137        // One tile's −Σ_set(x) from a TBL-unpacked pair load.
6138        macro_rules! tile_dot_tbl {
6139            ($ld:expr, $i0:expr, $i1:expr, $x:expr) => {{
6140                let w0 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8($ld, $i0), m));
6141                let w1 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8($ld, $i1), m));
6142                let x0 = vld1q_s8($x);
6143                let x1 = vld1q_s8($x.add(16));
6144                let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
6145                asm!(
6146                    "sdot {a0:v}.4s, {w0:v}.16b, {x0:v}.16b",
6147                    "sdot {a1:v}.4s, {w1:v}.16b, {x1:v}.16b",
6148                    a0 = inout(vreg) a0, a1 = inout(vreg) a1,
6149                    w0 = in(vreg) w0, x0 = in(vreg) x0, w1 = in(vreg) w1, x1 = in(vreg) x1,
6150                    options(pure, nomem, nostack),
6151                );
6152                vaddq_s32(a0, a1)
6153            }};
6154        }
6155        let base = bytes.as_ptr().add(r * gpr * Q1_TILE);
6156        let row_base = r * gpr * Q1_TILE;
6157        let abs_end = bytes.len();
6158        let xp = xq.as_ptr();
6159        let gp = gsum.as_ptr();
6160        let mut accv = vdupq_n_f32(0.0);
6161        let mut gi = 0;
6162        // The second pair load reads 4B past tile gi+3 — stay inside
6163        // the payload slice (only the file's final tiles fall back).
6164        while gi + 4 <= gpr && row_base + (gi + 4) * Q1_TILE + 4 <= abs_end {
6165            let t0 = base.add(gi * Q1_TILE);
6166            let ld_a = vld1q_u8(t0);
6167            let ld_b = vld1q_u8(t0.add(2 * Q1_TILE));
6168            let d0 = tile_dot_tbl!(ld_a, iw00, iw01, xp.add(gi * GROUP_SIZE));
6169            let d1 = tile_dot_tbl!(ld_a, iw10, iw11, xp.add((gi + 1) * GROUP_SIZE));
6170            let d2 = tile_dot_tbl!(ld_b, iw00, iw01, xp.add((gi + 2) * GROUP_SIZE));
6171            let d3 = tile_dot_tbl!(ld_b, iw10, iw11, xp.add((gi + 3) * GROUP_SIZE));
6172            // [−Σ0, −Σ1, −Σ2, −Σ3] → dots = −(2·Σset_neg + gsum)
6173            let neg = vpaddq_s32(vpaddq_s32(d0, d1), vpaddq_s32(d2, d3));
6174            let g = vld1q_s32(gp.add(gi));
6175            let dots = vnegq_s32(vaddq_s32(vshlq_n_s32::<1>(neg), g));
6176            let sc16 = vqtbl2_u8(uint8x16x2_t(ld_a, ld_b), isc);
6177            let scf: float32x4_t;
6178            asm!(
6179                "fcvtl {o:v}.4s, {i:v}.4h",
6180                o = out(vreg) scf, i = in(vreg) sc16,
6181                options(pure, nomem, nostack),
6182            );
6183            accv = vfmaq_f32(accv, vcvtq_f32_s32(dots), scf);
6184            gi += 4;
6185        }
6186        let mut acc = vaddvq_f32(accv);
6187        while gi < gpr {
6188            let t = base.add(gi * Q1_TILE);
6189            let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
6190            let d = vaddvq_s32(tile_dot!(t, xp.add(gi * GROUP_SIZE)));
6191            acc += (-(2 * d + *gp.add(gi))) as f32 * s;
6192            gi += 1;
6193        }
6194        acc
6195    }
6196}
6197
6198/// Blocked q1 1×4: one TBL unpack of the tile pair serves FOUR
6199/// activation streams (prefill amortization — the same idea as the
6200/// AVX2 twin; per stream the group order, fma order and tail match the
6201/// single-row kernel exactly, so batch == matvec bit-for-bit).
6202#[cfg(target_arch = "aarch64")]
6203#[target_feature(enable = "neon,dotprod")]
6204unsafe fn dot_q1_row_1x4_sdot(
6205    bytes: &[u8],
6206    r: usize,
6207    gpr: usize,
6208    xs: [&[i8]; 4],
6209    gs: [&[i32]; 4],
6210) -> [f32; 4] {
6211    // SAFETY: same slice-length contracts as `dot_q1_row_sdot`, ×4.
6212    unsafe {
6213        use core::arch::aarch64::*;
6214        use core::arch::asm;
6215        const MASKS: [u8; 16] = [1, 2, 4, 8, 16, 32, 64, 128, 1, 2, 4, 8, 16, 32, 64, 128];
6216        const IW00: [u8; 16] = [2, 2, 2, 2, 2, 2, 2, 2, 3, 3, 3, 3, 3, 3, 3, 3];
6217        const IW01: [u8; 16] = [4, 4, 4, 4, 4, 4, 4, 4, 5, 5, 5, 5, 5, 5, 5, 5];
6218        const IW10: [u8; 16] = [8, 8, 8, 8, 8, 8, 8, 8, 9, 9, 9, 9, 9, 9, 9, 9];
6219        const IW11: [u8; 16] = [
6220            10, 10, 10, 10, 10, 10, 10, 10, 11, 11, 11, 11, 11, 11, 11, 11,
6221        ];
6222        const ISC: [u8; 8] = [0, 1, 6, 7, 16, 17, 22, 23];
6223        let m = vld1q_u8(MASKS.as_ptr());
6224        let (iw00, iw01) = (vld1q_u8(IW00.as_ptr()), vld1q_u8(IW01.as_ptr()));
6225        let (iw10, iw11) = (vld1q_u8(IW10.as_ptr()), vld1q_u8(IW11.as_ptr()));
6226        let isc = vld1_u8(ISC.as_ptr());
6227        macro_rules! sdot2 {
6228            ($w0:expr, $w1:expr, $x:expr) => {{
6229                let x0 = vld1q_s8($x);
6230                let x1 = vld1q_s8($x.add(16));
6231                let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
6232                asm!(
6233                    "sdot {a0:v}.4s, {w0:v}.16b, {x0:v}.16b",
6234                    "sdot {a1:v}.4s, {w1:v}.16b, {x1:v}.16b",
6235                    a0 = inout(vreg) a0, a1 = inout(vreg) a1,
6236                    w0 = in(vreg) $w0, x0 = in(vreg) x0, w1 = in(vreg) $w1, x1 = in(vreg) x1,
6237                    options(pure, nomem, nostack),
6238                );
6239                vaddq_s32(a0, a1)
6240            }};
6241        }
6242        let base = bytes.as_ptr().add(r * gpr * Q1_TILE);
6243        let row_base = r * gpr * Q1_TILE;
6244        let abs_end = bytes.len();
6245        let mut accv = [vdupq_n_f32(0.0); 4];
6246        let mut gi = 0;
6247        while gi + 4 <= gpr && row_base + (gi + 4) * Q1_TILE + 4 <= abs_end {
6248            let t0 = base.add(gi * Q1_TILE);
6249            let ld_a = vld1q_u8(t0);
6250            let ld_b = vld1q_u8(t0.add(2 * Q1_TILE));
6251            // Unpack ONCE — eight ±mask vectors serve all four streams.
6252            let w00 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8(ld_a, iw00), m));
6253            let w01 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8(ld_a, iw01), m));
6254            let w10 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8(ld_a, iw10), m));
6255            let w11 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8(ld_a, iw11), m));
6256            let w20 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8(ld_b, iw00), m));
6257            let w21 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8(ld_b, iw01), m));
6258            let w30 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8(ld_b, iw10), m));
6259            let w31 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8(ld_b, iw11), m));
6260            let sc16 = vqtbl2_u8(uint8x16x2_t(ld_a, ld_b), isc);
6261            let scf: float32x4_t;
6262            asm!(
6263                "fcvtl {o:v}.4s, {i:v}.4h",
6264                o = out(vreg) scf, i = in(vreg) sc16,
6265                options(pure, nomem, nostack),
6266            );
6267            for k in 0..4 {
6268                let xp = xs[k].as_ptr();
6269                let d0 = sdot2!(w00, w01, xp.add(gi * GROUP_SIZE));
6270                let d1 = sdot2!(w10, w11, xp.add((gi + 1) * GROUP_SIZE));
6271                let d2 = sdot2!(w20, w21, xp.add((gi + 2) * GROUP_SIZE));
6272                let d3 = sdot2!(w30, w31, xp.add((gi + 3) * GROUP_SIZE));
6273                let neg = vpaddq_s32(vpaddq_s32(d0, d1), vpaddq_s32(d2, d3));
6274                let g = vld1q_s32(gs[k].as_ptr().add(gi));
6275                let dots = vnegq_s32(vaddq_s32(vshlq_n_s32::<1>(neg), g));
6276                accv[k] = vfmaq_f32(accv[k], vcvtq_f32_s32(dots), scf);
6277            }
6278            gi += 4;
6279        }
6280        let mut acc = [
6281            vaddvq_f32(accv[0]),
6282            vaddvq_f32(accv[1]),
6283            vaddvq_f32(accv[2]),
6284            vaddvq_f32(accv[3]),
6285        ];
6286        while gi < gpr {
6287            let t = base.add(gi * Q1_TILE);
6288            let sc = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
6289            let v0 = vcombine_u8(vdup_n_u8(*t.add(2)), vdup_n_u8(*t.add(3)));
6290            let v1 = vcombine_u8(vdup_n_u8(*t.add(4)), vdup_n_u8(*t.add(5)));
6291            let w0 = vreinterpretq_s8_u8(vtstq_u8(v0, m));
6292            let w1 = vreinterpretq_s8_u8(vtstq_u8(v1, m));
6293            for k in 0..4 {
6294                let d = vaddvq_s32(sdot2!(w0, w1, xs[k].as_ptr().add(gi * GROUP_SIZE)));
6295                acc[k] += (-(2 * d + *gs[k].as_ptr().add(gi))) as f32 * sc;
6296            }
6297            gi += 1;
6298        }
6299        acc
6300    }
6301}
6302
6303/// (weight ±1, scale) of one q1 element — the exact outlier term.
6304#[inline]
6305fn q1_outlier(bytes: &[u8], r: usize, gpr: usize, j: usize) -> (f32, f32) {
6306    let gi = j / GROUP_SIZE;
6307    let k = j % GROUP_SIZE;
6308    let tile = &bytes[(r * gpr + gi) * Q1_TILE..(r * gpr + gi + 1) * Q1_TILE];
6309    let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
6310    let bit = (tile[2 + k / 8] >> (k % 8)) & 1;
6311    ((bit as i32 * 2 - 1) as f32, s)
6312}
6313
6314/// Exact scalar q1 row (CMF_SDOT=0 contract).
6315#[inline]
6316fn q1_row_exact(bytes: &[u8], r: usize, gpr: usize, x: &[f32]) -> f32 {
6317    let mut acc = 0f32;
6318    for gi in 0..gpr {
6319        let tile = &bytes[(r * gpr + gi) * Q1_TILE..(r * gpr + gi + 1) * Q1_TILE];
6320        let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
6321        let xg = &x[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
6322        let mut ga = 0f32;
6323        for (j, &b) in tile[2..].iter().enumerate() {
6324            for k in 0..8 {
6325                ga += (((b >> k) & 1) as f32 * 2.0 - 1.0) * xg[j * 8 + k];
6326            }
6327        }
6328        acc += ga * s;
6329    }
6330    acc
6331}
6332
6333/// One q1 row range via A8W8 (the body of `q1_matvec`'s hot loop,
6334/// extracted so multi-matrix jobs drive the same kernel).
6335#[allow(clippy::too_many_arguments)]
6336fn q1_range_a8w8(
6337    bytes: &[u8],
6338    gpr: usize,
6339    act: &SplitAct,
6340    gsum: &[i32],
6341    out: SendMut,
6342    start: usize,
6343    end: usize,
6344) {
6345    for r in start..end {
6346        let mut acc = dot_q1_row_i8(bytes, r, gpr, &act.xq, gsum) * act.sx;
6347        for &(j, xv) in &act.outliers {
6348            let (w, s) = q1_outlier(bytes, r, gpr, j);
6349            acc += w * s * xv;
6350        }
6351        // SAFETY: disjoint row ranges per worker.
6352        unsafe { *out.at(r) = acc };
6353    }
6354}
6355
6356/// Exact-scalar q1 row range (CMF_SDOT=0 contract).
6357fn q1_range_f32(bytes: &[u8], gpr: usize, x: &[f32], out: SendMut, start: usize, end: usize) {
6358    for r in start..end {
6359        // SAFETY: disjoint row ranges per worker.
6360        unsafe { *out.at(r) = q1_row_exact(bytes, r, gpr, x) };
6361    }
6362}
6363
6364/// q1t per-row overlay locator. After the base (`base_len`) come
6365/// `[u32 row_ptr[rows+1]]` then `[(u16 col, f16 val)]` grouped by row (row
6366/// `r`'s entries are `[row_ptr[r], row_ptr[r+1])`). Returns
6367/// `(row_ptr offset, entries offset, present)`.
6368fn q1t_overlay(bytes: &[u8], base_len: usize, rows: usize) -> (usize, usize, bool) {
6369    let entries = base_len + (rows + 1) * 4;
6370    (base_len, entries, entries <= bytes.len())
6371}
6372
6373/// Read `row_ptr[r]` from the overlay's prefix-sum table.
6374#[inline]
6375fn q1t_rowptr(bytes: &[u8], rp_off: usize, r: usize) -> usize {
6376    let o = rp_off + r * 4;
6377    u32::from_le_bytes([bytes[o], bytes[o + 1], bytes[o + 2], bytes[o + 3]]) as usize
6378}
6379
6380/// Byte → the 5 ternary signs it packs `{−1,0,+1}` as f32, precomputed so
6381/// decoding a q1t code is a table load, not the base-3 divide/modulo per
6382/// weight (division is ~20–40× the cost of a load). Built at compile time.
6383const SIGN5: [[f32; 5]; 256] = {
6384    let mut lut = [[0.0f32; 5]; 256];
6385    let pow3 = [1u16, 3, 9, 27, 81];
6386    let mut byte = 0usize;
6387    while byte < 256 {
6388        let mut i = 0usize;
6389        while i < 5 {
6390            let code = (byte as u16 / pow3[i]) % 3;
6391            lut[byte][i] = if code == 1 {
6392                1.0
6393            } else if code == 2 {
6394                -1.0
6395            } else {
6396                0.0
6397            };
6398            i += 1;
6399        }
6400        byte += 1;
6401    }
6402    lut
6403};
6404
6405/// Same table, as i8 signs — the operand for the int8 SDOT base kernel.
6406const SIGN5_I8: [[i8; 5]; 256] = {
6407    let mut lut = [[0i8; 5]; 256];
6408    let pow3 = [1u16, 3, 9, 27, 81];
6409    let mut byte = 0usize;
6410    while byte < 256 {
6411        let mut i = 0usize;
6412        while i < 5 {
6413            let code = (byte as u16 / pow3[i]) % 3;
6414            lut[byte][i] = if code == 1 {
6415                1
6416            } else if code == 2 {
6417                -1
6418            } else {
6419                0
6420            };
6421            i += 1;
6422        }
6423        byte += 1;
6424    }
6425    lut
6426};
6427
6428/// The same 5 i8 signs packed into a u64 (`[s0 s1 s2 s3 s4 0 0 0]`, LE) so the
6429/// group unpack is 7 unaligned u64 stores at offsets 0,5,10,…,30 instead of
6430/// six 5-byte copies + LUT indexing — each store's trailing zeros are fixed by
6431/// the next store, and the last one runs 6 B past the 32nd weight (the unpack
6432/// buffer is padded to 40). This is the decode/prefill hot inner op.
6433const SIGN5_U64: [u64; 256] = {
6434    let mut lut = [0u64; 256];
6435    let pow3 = [1u16, 3, 9, 27, 81];
6436    let mut byte = 0usize;
6437    while byte < 256 {
6438        let mut v = 0u64;
6439        let mut i = 0usize;
6440        while i < 5 {
6441            let code = (byte as u16 / pow3[i]) % 3;
6442            let s: u8 = if code == 1 {
6443                1
6444            } else if code == 2 {
6445                0xFF
6446            } else {
6447                0
6448            };
6449            v |= (s as u64) << (i * 8);
6450            i += 1;
6451        }
6452        lut[byte] = v;
6453        byte += 1;
6454    }
6455    lut
6456};
6457
6458/// Ternary base weight at `(row r, col j)` = `sign(code)·s_group`. Used to add
6459/// back activation-outlier columns, whose `x` was zeroed for the int8 bulk dot
6460/// (`split_act`). At a weight-outlier position the code is 0, so this is 0 and
6461/// the overlay correction owns that column — no double counting.
6462#[inline]
6463fn q1t_base_weight(bytes: &[u8], r: usize, gpr: usize, j: usize) -> f32 {
6464    const TILE: usize = cortiq_core::quant::Q1T_TILE;
6465    let off = (r * gpr + j / GROUP_SIZE) * TILE;
6466    let s = f16_to_f32(u16::from_le_bytes([bytes[off], bytes[off + 1]]));
6467    let within = j % GROUP_SIZE;
6468    SIGN5[bytes[off + 2 + within / 5] as usize][within % 5] * s
6469}
6470
6471/// One 32-group int8 dot via two SDOTs. Bit-exact vs the scalar i8 sum
6472/// (integer accumulation is order-independent).
6473#[cfg(target_arch = "aarch64")]
6474#[target_feature(enable = "neon,dotprod")]
6475#[inline]
6476unsafe fn sdot32_i8(w: *const i8, x: *const i8) -> i32 {
6477    // SAFETY: caller guarantees 32 readable i8 at each pointer.
6478    unsafe {
6479        use core::arch::aarch64::*;
6480        use core::arch::asm;
6481        let w0 = vld1q_s8(w);
6482        let w1 = vld1q_s8(w.add(16));
6483        let x0 = vld1q_s8(x);
6484        let x1 = vld1q_s8(x.add(16));
6485        let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
6486        asm!(
6487            "sdot {a0:v}.4s, {w0:v}.16b, {x0:v}.16b",
6488            "sdot {a1:v}.4s, {w1:v}.16b, {x1:v}.16b",
6489            a0 = inout(vreg) a0, a1 = inout(vreg) a1,
6490            w0 = in(vreg) w0, x0 = in(vreg) x0, w1 = in(vreg) w1, x1 = in(vreg) x1,
6491            options(pure, nomem, nostack),
6492        );
6493        vaddvq_s32(vaddq_s32(a0, a1))
6494    }
6495}
6496
6497/// One 32-group int8 dot via AVX2: signed·signed as `maddubs(|w|, sign(x,w))`
6498/// then `madd` and a horizontal reduce (the same idiom as `dot_q4t_row_avx2`).
6499#[cfg(target_arch = "x86_64")]
6500#[target_feature(enable = "avx2")]
6501#[inline]
6502unsafe fn i8dot32_avx2(w: *const i8, x: *const i8) -> i32 {
6503    // SAFETY: caller guarantees 32 readable i8 at each pointer.
6504    unsafe {
6505        use core::arch::x86_64::*;
6506        let wv = _mm256_loadu_si256(w as *const __m256i);
6507        let xv = _mm256_loadu_si256(x as *const __m256i);
6508        let p16 = _mm256_maddubs_epi16(_mm256_abs_epi8(wv), _mm256_sign_epi8(xv, wv));
6509        let d = _mm256_madd_epi16(p16, _mm256_set1_epi16(1));
6510        let hi128 = _mm256_extracti128_si256::<1>(d);
6511        let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
6512        let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
6513        let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
6514        _mm_cvtsi128_si32(s32)
6515    }
6516}
6517
6518/// Unpack one q1t group's base-3 codes into 32 i8 signs via 7 unaligned u64
6519/// stores (see `SIGN5_U64`). `dst` MUST have ≥ 40 bytes: the 7th store writes
6520/// `dst[30..38]`. Stores go in order so each one's trailing zeros are
6521/// overwritten by the next; the final 6 padding bytes are unused by the dot.
6522#[inline]
6523fn q1t_unpack_group_i8(codes: *const u8, dst: &mut [i8]) {
6524    debug_assert!(dst.len() >= 40);
6525    // SAFETY: codes points at 7 readable bytes; dst has ≥ 40 bytes so every
6526    // 8-byte store at offset bi*5 (bi ≤ 6 → ≤ 30) stays in bounds.
6527    unsafe {
6528        let p = dst.as_mut_ptr();
6529        for bi in 0..7 {
6530            core::ptr::write_unaligned(
6531                p.add(bi * 5) as *mut u64,
6532                SIGN5_U64[*codes.add(bi) as usize],
6533            );
6534        }
6535    }
6536}
6537
6538/// One 32-group int8 dot, arch-dispatched (the matmat inner loop, where the
6539/// row's signs are unpacked once and dotted against every batch input).
6540/// Callers are gated by `a8w8_enabled()`, so the target-feature arms are
6541/// reachable; the scalar arm is a non-SIMD-arch fallback.
6542#[inline]
6543fn q1t_i8dot32(w: *const i8, x: *const i8) -> i32 {
6544    #[cfg(target_arch = "aarch64")]
6545    unsafe {
6546        return sdot32_i8(w, x);
6547    }
6548    #[cfg(target_arch = "x86_64")]
6549    unsafe {
6550        return i8dot32_avx2(w, x);
6551    }
6552    #[allow(unreachable_code)]
6553    unsafe {
6554        let mut s = 0i32;
6555        for k in 0..GROUP_SIZE {
6556            s += *w.add(k) as i32 * *x.add(k) as i32;
6557        }
6558        s
6559    }
6560}
6561
6562#[inline]
6563unsafe fn q1t_unpack_reg_u64s(codes: *const u8) -> (u64, u64, u64, u64) {
6564    let (s0, s1, s2, s3, s4, s5, s6) = unsafe {
6565        (
6566            SIGN5_U64[*codes as usize],
6567            SIGN5_U64[*codes.add(1) as usize],
6568            SIGN5_U64[*codes.add(2) as usize],
6569            SIGN5_U64[*codes.add(3) as usize],
6570            SIGN5_U64[*codes.add(4) as usize],
6571            SIGN5_U64[*codes.add(5) as usize],
6572            SIGN5_U64[*codes.add(6) as usize],
6573        )
6574    };
6575
6576    let u0 = s0 | (s1 << 40);
6577    let u1 = (s1 >> 24) | (s2 << 16) | (s3 << 56);
6578    let u2 = (s3 >> 8) | (s4 << 32);
6579    let u3 = (s4 >> 32) | (s5 << 8) | (s6 << 48);
6580
6581    (u0, u1, u2, u3)
6582}
6583
6584/// One q1t row's int8 base dot: `Σ_group s·dot(signs, xq)` (before the shared
6585/// `sx`). Direct register unpacking (zero stack stores/loads, no STLF stalls).
6586/// ARM SDOT.
6587#[cfg(target_arch = "aarch64")]
6588#[target_feature(enable = "neon,dotprod")]
6589unsafe fn q1t_dot_row_sdot(bytes: &[u8], r: usize, gpr: usize, xq: &[i8]) -> f32 {
6590    use core::arch::aarch64::*;
6591    use core::arch::asm;
6592    unsafe {
6593        const TILE: usize = cortiq_core::quant::Q1T_TILE;
6594        let mut acc = 0f32;
6595        let bytes_ptr = bytes.as_ptr();
6596        let xq_ptr = xq.as_ptr();
6597        let row_off = r * gpr * TILE;
6598
6599        let gpr2 = gpr & !1;
6600        let mut gi = 0;
6601        while gi < gpr2 {
6602            let off0 = row_off + gi * TILE;
6603            let off1 = off0 + TILE;
6604            let s0 = f16_to_f32(u16::from_le_bytes([
6605                *bytes_ptr.add(off0),
6606                *bytes_ptr.add(off0 + 1),
6607            ]));
6608            let s1 = f16_to_f32(u16::from_le_bytes([
6609                *bytes_ptr.add(off1),
6610                *bytes_ptr.add(off1 + 1),
6611            ]));
6612
6613            let (u0_0, u1_0, u2_0, u3_0) = q1t_unpack_reg_u64s(bytes_ptr.add(off0 + 2));
6614            let (u0_1, u1_1, u2_1, u3_1) = q1t_unpack_reg_u64s(bytes_ptr.add(off1 + 2));
6615
6616            let w0_0 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u0_0), vcreate_u64(u1_0)));
6617            let w1_0 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u2_0), vcreate_u64(u3_0)));
6618            let w0_1 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u0_1), vcreate_u64(u1_1)));
6619            let w1_1 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u2_1), vcreate_u64(u3_1)));
6620
6621            let x0_0 = vld1q_s8(xq_ptr.add(gi * GROUP_SIZE));
6622            let x1_0 = vld1q_s8(xq_ptr.add(gi * GROUP_SIZE + 16));
6623            let x0_1 = vld1q_s8(xq_ptr.add((gi + 1) * GROUP_SIZE));
6624            let x1_1 = vld1q_s8(xq_ptr.add((gi + 1) * GROUP_SIZE + 16));
6625
6626            let (mut a0_0, mut a1_0) = (vdupq_n_s32(0), vdupq_n_s32(0));
6627            let (mut a0_1, mut a1_1) = (vdupq_n_s32(0), vdupq_n_s32(0));
6628            asm!(
6629                "sdot {a0_0:v}.4s, {w0_0:v}.16b, {x0_0:v}.16b",
6630                "sdot {a1_0:v}.4s, {w1_0:v}.16b, {x1_0:v}.16b",
6631                "sdot {a0_1:v}.4s, {w0_1:v}.16b, {x0_1:v}.16b",
6632                "sdot {a1_1:v}.4s, {w1_1:v}.16b, {x1_1:v}.16b",
6633                a0_0 = inout(vreg) a0_0, a1_0 = inout(vreg) a1_0,
6634                a0_1 = inout(vreg) a0_1, a1_1 = inout(vreg) a1_1,
6635                w0_0 = in(vreg) w0_0, x0_0 = in(vreg) x0_0, w1_0 = in(vreg) w1_0, x1_0 = in(vreg) x1_0,
6636                w0_1 = in(vreg) w0_1, x0_1 = in(vreg) x0_1, w1_1 = in(vreg) w1_1, x1_1 = in(vreg) x1_1,
6637                options(pure, nomem, nostack),
6638            );
6639            let d0 = vaddvq_s32(vaddq_s32(a0_0, a1_0));
6640            let d1 = vaddvq_s32(vaddq_s32(a0_1, a1_1));
6641            acc += d0 as f32 * s0 + d1 as f32 * s1;
6642            gi += 2;
6643        }
6644
6645        if gi < gpr {
6646            let off = row_off + gi * TILE;
6647            let s = f16_to_f32(u16::from_le_bytes([
6648                *bytes_ptr.add(off),
6649                *bytes_ptr.add(off + 1),
6650            ]));
6651            let (u0, u1, u2, u3) = q1t_unpack_reg_u64s(bytes_ptr.add(off + 2));
6652            let w0 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u0), vcreate_u64(u1)));
6653            let w1 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u2), vcreate_u64(u3)));
6654            let x0 = vld1q_s8(xq_ptr.add(gi * GROUP_SIZE));
6655            let x1 = vld1q_s8(xq_ptr.add(gi * GROUP_SIZE + 16));
6656            let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
6657            asm!(
6658                "sdot {a0:v}.4s, {w0:v}.16b, {x0:v}.16b",
6659                "sdot {a1:v}.4s, {w1:v}.16b, {x1:v}.16b",
6660                a0 = inout(vreg) a0, a1 = inout(vreg) a1,
6661                w0 = in(vreg) w0, x0 = in(vreg) x0, w1 = in(vreg) w1, x1 = in(vreg) x1,
6662                options(pure, nomem, nostack),
6663            );
6664            let d = vaddvq_s32(vaddq_s32(a0, a1));
6665            acc += d as f32 * s;
6666        }
6667        acc
6668    }
6669}
6670
6671/// x86 AVX2 mirror of `q1t_dot_row_sdot` (maddubs int8 dot per group).
6672#[cfg(target_arch = "x86_64")]
6673#[target_feature(enable = "avx2")]
6674unsafe fn q1t_dot_row_avx2(bytes: &[u8], r: usize, gpr: usize, xq: &[i8]) -> f32 {
6675    use core::arch::x86_64::*;
6676    unsafe {
6677        const TILE: usize = cortiq_core::quant::Q1T_TILE;
6678        let mut acc = 0f32;
6679        let bytes_ptr = bytes.as_ptr();
6680        let xq_ptr = xq.as_ptr();
6681        let row_off = r * gpr * TILE;
6682
6683        let ones = _mm256_set1_epi16(1);
6684        for gi in 0..gpr {
6685            let off = row_off + gi * TILE;
6686            let s = f16_to_f32(u16::from_le_bytes([
6687                *bytes_ptr.add(off),
6688                *bytes_ptr.add(off + 1),
6689            ]));
6690            let (u0, u1, u2, u3) = q1t_unpack_reg_u64s(bytes_ptr.add(off + 2));
6691            let wv = _mm256_set_epi64x(u3 as i64, u2 as i64, u1 as i64, u0 as i64);
6692            let xv = _mm256_loadu_si256(xq_ptr.add(gi * GROUP_SIZE) as *const __m256i);
6693            let p16 = _mm256_maddubs_epi16(_mm256_abs_epi8(wv), _mm256_sign_epi8(xv, wv));
6694            let d256 = _mm256_madd_epi16(p16, ones);
6695            let d128 = _mm_add_epi32(
6696                _mm256_castsi256_si128(d256),
6697                _mm256_extracti128_si256(d256, 1),
6698            );
6699            let d64 = _mm_add_epi32(d128, _mm_shuffle_epi32(d128, 0xee));
6700            let d32 = _mm_cvtsi128_si32(_mm_add_epi32(d64, _mm_shuffle_epi32(d64, 0x55)));
6701            acc += d32 as f32 * s;
6702        }
6703        acc
6704    }
6705}
6706
6707/// VNNI twin of `q1t_dot_row_avx2` (see `dpbusd_hsum`).
6708#[cfg(target_arch = "x86_64")]
6709#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
6710unsafe fn q1t_dot_row_vnni(bytes: &[u8], r: usize, gpr: usize, xq: &[i8]) -> f32 {
6711    use core::arch::x86_64::*;
6712    // SAFETY: same tile/xq contracts as `q1t_dot_row_avx2`.
6713    unsafe {
6714        const TILE: usize = cortiq_core::quant::Q1T_TILE;
6715        let mut acc = 0f32;
6716        let bytes_ptr = bytes.as_ptr();
6717        let xq_ptr = xq.as_ptr();
6718        let row_off = r * gpr * TILE;
6719        for gi in 0..gpr {
6720            let off = row_off + gi * TILE;
6721            let s = f16_to_f32(u16::from_le_bytes([
6722                *bytes_ptr.add(off),
6723                *bytes_ptr.add(off + 1),
6724            ]));
6725            let (u0, u1, u2, u3) = q1t_unpack_reg_u64s(bytes_ptr.add(off + 2));
6726            let wv = _mm256_set_epi64x(u3 as i64, u2 as i64, u1 as i64, u0 as i64);
6727            let xv = _mm256_loadu_si256(xq_ptr.add(gi * GROUP_SIZE) as *const __m256i);
6728            let d = dpbusd_hsum(_mm256_abs_epi8(wv), _mm256_sign_epi8(xv, wv));
6729            acc += d as f32 * s;
6730        }
6731        acc
6732    }
6733}
6734
6735/// Per-row int8 base dot, dispatched once per row (matvec decode hot path).
6736/// Callers are gated by `a8w8_enabled()`, so the target-feature kernels are
6737/// reachable.
6738#[inline]
6739fn q1t_dot_row_i8(bytes: &[u8], r: usize, gpr: usize, xq: &[i8]) -> f32 {
6740    #[cfg(target_arch = "aarch64")]
6741    unsafe {
6742        return q1t_dot_row_sdot(bytes, r, gpr, xq);
6743    }
6744    #[cfg(target_arch = "x86_64")]
6745    unsafe {
6746        if vnni_tiles_enabled() {
6747            return q1t_dot_row_vnni(bytes, r, gpr, xq);
6748        }
6749        return q1t_dot_row_avx2(bytes, r, gpr, xq);
6750    }
6751    #[allow(unreachable_code)]
6752    {
6753        const TILE: usize = cortiq_core::quant::Q1T_TILE;
6754        let mut acc = 0f32;
6755        let mut sg = [0i8; GROUP_SIZE + 8]; // +8 slack for the u64-store unpack
6756        for gi in 0..gpr {
6757            let off = (r * gpr + gi) * TILE;
6758            let s = f16_to_f32(u16::from_le_bytes([bytes[off], bytes[off + 1]]));
6759            q1t_unpack_group_i8(bytes.as_ptr().wrapping_add(off + 2), &mut sg);
6760            let mut d = 0i32;
6761            for k in 0..GROUP_SIZE {
6762                d += sg[k] as i32 * xq[gi * GROUP_SIZE + k] as i32;
6763            }
6764            acc += d as f32 * s;
6765        }
6766        acc
6767    }
6768}
6769
6770/// Σ over a row's outliers of `value·x[col]` — the correction that adds the
6771/// overlay's exact weights on top of the base dot. INVARIANT: the encoder
6772/// writes ternary code 0 at every outlier position (`quantize_q1t`), so the
6773/// base contributes nothing there and this is a plain `value·x`, not
6774/// `(value − base)·x` — no scattered per-outlier scale read. Row `r`'s entries
6775/// are the contiguous slice `[row_ptr[r], row_ptr[r+1])`, so no binary search.
6776fn q1t_row_outlier_correction(
6777    bytes: &[u8],
6778    r: usize,
6779    rp_off: usize,
6780    entries_off: usize,
6781    has_ov: bool,
6782    x: &[f32],
6783) -> f32 {
6784    if !has_ov {
6785        return 0.0;
6786    }
6787    let (c0, c1) = (
6788        q1t_rowptr(bytes, rp_off, r),
6789        q1t_rowptr(bytes, rp_off, r + 1),
6790    );
6791    let mut corr = 0f32;
6792    for p in c0..c1 {
6793        let e = entries_off + p * 4;
6794        let col = u16::from_le_bytes([bytes[e], bytes[e + 1]]) as usize;
6795        let val = f16_to_f32(u16::from_le_bytes([bytes[e + 2], bytes[e + 3]]));
6796        corr += val * x[col];
6797    }
6798    corr
6799}
6800
6801/// Dequantize one q1t row into `buf[..cols]` via the sign LUT (no division),
6802/// then apply the row's outliers (its `[row_ptr[r], row_ptr[r+1])` slice).
6803/// Used by the batched (prefill) path where the decode amortizes over the batch.
6804fn q1t_dequant_row(
6805    bytes: &[u8],
6806    r: usize,
6807    gpr: usize,
6808    rp_off: usize,
6809    entries_off: usize,
6810    has_ov: bool,
6811    buf: &mut [f32],
6812) {
6813    const TILE: usize = cortiq_core::quant::Q1T_TILE;
6814    for g in 0..gpr {
6815        let off = (r * gpr + g) * TILE;
6816        let s = f16_to_f32(u16::from_le_bytes([bytes[off], bytes[off + 1]]));
6817        let codes = &bytes[off + 2..off + TILE];
6818        let bc = g * GROUP_SIZE;
6819        // 6 full bytes (30 codes) + a 7th byte holding the last 2.
6820        for bi in 0..6 {
6821            let lut = &SIGN5[codes[bi] as usize];
6822            let d = &mut buf[bc + bi * 5..bc + bi * 5 + 5];
6823            for i in 0..5 {
6824                d[i] = lut[i] * s;
6825            }
6826        }
6827        let lut = &SIGN5[codes[6] as usize];
6828        buf[bc + 30] = lut[0] * s;
6829        buf[bc + 31] = lut[1] * s;
6830    }
6831    if !has_ov {
6832        return;
6833    }
6834    let (c0, c1) = (
6835        q1t_rowptr(bytes, rp_off, r),
6836        q1t_rowptr(bytes, rp_off, r + 1),
6837    );
6838    for p in c0..c1 {
6839        let e = entries_off + p * 4;
6840        let col = u16::from_le_bytes([bytes[e], bytes[e + 1]]) as usize;
6841        buf[col] = f16_to_f32(u16::from_le_bytes([bytes[e + 2], bytes[e + 3]]));
6842    }
6843}
6844
6845/// Add the sparse outlier overlay onto a base dot already in `out` (the GPU
6846/// computes the ternary base; the overlay stays on the CPU — its entries are
6847/// few and its per-row gather doesn't vectorize on the GPU). Row-parallel.
6848fn q1t_add_overlay(
6849    bytes: &[u8],
6850    x: &[f32],
6851    rows: usize,
6852    cols: usize,
6853    out: &mut [f32],
6854    pool: Option<&Pool>,
6855) {
6856    const TILE: usize = cortiq_core::quant::Q1T_TILE;
6857    let gpr = cols / GROUP_SIZE;
6858    let (rp_off, ent_off, has_ov) = q1t_overlay(bytes, rows * gpr * TILE, rows);
6859    if !has_ov {
6860        return;
6861    }
6862    let out_addr = SendMut(out.as_mut_ptr());
6863    let run = move |start: usize, end: usize| {
6864        for r in start..end {
6865            let corr = q1t_row_outlier_correction(bytes, r, rp_off, ent_off, has_ov, x);
6866            // SAFETY: disjoint rows; add onto the base the GPU already wrote.
6867            unsafe { *out_addr.at(r) += corr };
6868        }
6869    };
6870    dispatch_rows(pool, rows, &run);
6871}
6872
6873/// Q1T row range via the A8W8 int8 path — shared activation split,
6874/// per-row: base SDOT dot + outlier correction + overlay.
6875#[allow(clippy::too_many_arguments)]
6876fn q1t_range_a8w8(
6877    bytes: &[u8],
6878    gpr: usize,
6879    rp_off: usize,
6880    ent_off: usize,
6881    has_ov: bool,
6882    act: &SplitAct,
6883    x: &[f32],
6884    out: SendMut,
6885    start: usize,
6886    end: usize,
6887) {
6888    for r in start..end {
6889        let mut acc = q1t_dot_row_i8(bytes, r, gpr, &act.xq) * act.sx;
6890        for &(j, xv) in &act.outliers {
6891            acc += q1t_base_weight(bytes, r, gpr, j) * xv;
6892        }
6893        acc += q1t_row_outlier_correction(bytes, r, rp_off, ent_off, has_ov, x);
6894        // SAFETY: disjoint row ranges per worker.
6895        unsafe { *out.at(r) = acc };
6896    }
6897}
6898
6899/// Q1T row range via the f32 path (no SDOT) — for matvec_many batched
6900/// dispatch when a8w8 is unavailable.
6901#[allow(clippy::too_many_arguments)]
6902fn q1t_range_f32_batch(
6903    bytes: &[u8],
6904    gpr: usize,
6905    rp_off: usize,
6906    ent_off: usize,
6907    has_ov: bool,
6908    x: &[f32],
6909    out: SendMut,
6910    start: usize,
6911    end: usize,
6912) {
6913    const TILE: usize = cortiq_core::quant::Q1T_TILE;
6914    let mut sg = [0f32; GROUP_SIZE];
6915    for r in start..end {
6916        let mut acc = 0f32;
6917        for g in 0..gpr {
6918            let off = (r * gpr + g) * TILE;
6919            let s = f16_to_f32(u16::from_le_bytes([bytes[off], bytes[off + 1]]));
6920            let codes = &bytes[off + 2..off + TILE];
6921            let xg = &x[g * GROUP_SIZE..g * GROUP_SIZE + GROUP_SIZE];
6922            for bi in 0..6 {
6923                sg[bi * 5..bi * 5 + 5].copy_from_slice(&SIGN5[codes[bi] as usize]);
6924            }
6925            let lut = &SIGN5[codes[6] as usize];
6926            sg[30] = lut[0];
6927            sg[31] = lut[1];
6928            let mut gsum = 0f32;
6929            for k in 0..GROUP_SIZE {
6930                gsum += sg[k] * xg[k];
6931            }
6932            acc += s * gsum;
6933        }
6934        acc += q1t_row_outlier_correction(bytes, r, rp_off, ent_off, has_ov, x);
6935        // SAFETY: disjoint row ranges per worker.
6936        unsafe { *out.at(r) = acc };
6937    }
6938}
6939
6940/// Ternary (q1t) matvec — decode+dot straight from mmap, one group at a time:
6941/// no per-ROW buffer, no division (the sign LUT), and a tiny per-group sign
6942/// buffer so the 32-wide dot vectorizes. This is the decode hot path.
6943fn q1t_matvec(
6944    bytes: &[u8],
6945    x: &[f32],
6946    rows: usize,
6947    cols: usize,
6948    out: &mut [f32],
6949    pool: Option<&Pool>,
6950) {
6951    debug_assert_eq!(out.len(), rows);
6952    const TILE: usize = cortiq_core::quant::Q1T_TILE;
6953    let gpr = cols / GROUP_SIZE;
6954    let (rp_off, ent_off, has_ov) = q1t_overlay(bytes, rows * gpr * TILE, rows);
6955    let out_addr = SendMut(out.as_mut_ptr());
6956    // int8 SDOT base dot (ARM dotprod): ~4× the f32 arithmetic. x → i8 once
6957    // (`split_act`), activation outliers added back exactly in f32, weight
6958    // overlay on top. ARM SDOT / x86 AVX2; CMF_SDOT=0 keeps the exact f32 path.
6959    if a8w8_enabled() {
6960        let act = split_act(x);
6961        let act = &act;
6962        let run = move |start: usize, end: usize| {
6963            for r in start..end {
6964                let mut acc = q1t_dot_row_i8(bytes, r, gpr, &act.xq) * act.sx;
6965                for &(j, xv) in &act.outliers {
6966                    acc += q1t_base_weight(bytes, r, gpr, j) * xv;
6967                }
6968                acc += q1t_row_outlier_correction(bytes, r, rp_off, ent_off, has_ov, x);
6969                // SAFETY: disjoint row ranges per worker.
6970                unsafe { *out_addr.at(r) = acc };
6971            }
6972        };
6973        dispatch_rows(pool, rows, &run);
6974        return;
6975    }
6976    let run = move |start: usize, end: usize| {
6977        // Per-group signs, unpacked contiguously so the dot below is a clean
6978        // 32-wide reduction the autovectorizer turns into f32x4 FMAs — the
6979        // 5-values-per-byte base-3 layout won't SIMD in place.
6980        let mut sg = [0f32; GROUP_SIZE];
6981        for r in start..end {
6982            let mut acc = 0f32;
6983            for g in 0..gpr {
6984                let off = (r * gpr + g) * TILE;
6985                let s = f16_to_f32(u16::from_le_bytes([bytes[off], bytes[off + 1]]));
6986                let codes = &bytes[off + 2..off + TILE];
6987                let xg = &x[g * GROUP_SIZE..g * GROUP_SIZE + GROUP_SIZE];
6988                for bi in 0..6 {
6989                    sg[bi * 5..bi * 5 + 5].copy_from_slice(&SIGN5[codes[bi] as usize]);
6990                }
6991                let lut = &SIGN5[codes[6] as usize];
6992                sg[30] = lut[0];
6993                sg[31] = lut[1];
6994                let mut gsum = 0f32;
6995                for k in 0..GROUP_SIZE {
6996                    gsum += sg[k] * xg[k];
6997                }
6998                acc += s * gsum;
6999            }
7000            acc += q1t_row_outlier_correction(bytes, r, rp_off, ent_off, has_ov, x);
7001            unsafe { *out_addr.at(r) = acc };
7002        }
7003    };
7004    dispatch_rows(pool, rows, &run);
7005}
7006
7007/// Fused-pair twin of `q1t_dot_row_sdot`: ONE register unpack of the
7008/// ternary codes serves BOTH activation streams (the unpack chain is
7009/// the dominant per-row cost — MTP verify pairs paid it twice). Per
7010/// stream the group order and f32 accumulation match the single-row
7011/// kernel exactly, so pair == 2×matvec bit-for-bit.
7012#[cfg(target_arch = "aarch64")]
7013#[target_feature(enable = "neon,dotprod")]
7014unsafe fn q1t_dot_row_sdot2(bytes: &[u8], r: usize, gpr: usize, xa: &[i8], xb: &[i8]) -> [f32; 2] {
7015    use core::arch::aarch64::*;
7016    use core::arch::asm;
7017    // SAFETY: same slice-length contracts as `q1t_dot_row_sdot`, ×2.
7018    unsafe {
7019        const TILE: usize = cortiq_core::quant::Q1T_TILE;
7020        let bytes_ptr = bytes.as_ptr();
7021        let row_off = r * gpr * TILE;
7022        let xp = [xa.as_ptr(), xb.as_ptr()];
7023        let mut acc = [0f32; 2];
7024        macro_rules! sdot2 {
7025            ($w0:expr, $w1:expr, $x:expr) => {{
7026                let x0 = vld1q_s8($x);
7027                let x1 = vld1q_s8($x.add(16));
7028                let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
7029                asm!(
7030                    "sdot {a0:v}.4s, {w0:v}.16b, {x0:v}.16b",
7031                    "sdot {a1:v}.4s, {w1:v}.16b, {x1:v}.16b",
7032                    a0 = inout(vreg) a0, a1 = inout(vreg) a1,
7033                    w0 = in(vreg) $w0, x0 = in(vreg) x0, w1 = in(vreg) $w1, x1 = in(vreg) x1,
7034                    options(pure, nomem, nostack),
7035                );
7036                vaddvq_s32(vaddq_s32(a0, a1))
7037            }};
7038        }
7039        let gpr2 = gpr & !1;
7040        let mut gi = 0;
7041        while gi < gpr2 {
7042            let off0 = row_off + gi * TILE;
7043            let off1 = off0 + TILE;
7044            let s0 = f16_to_f32(u16::from_le_bytes([
7045                *bytes_ptr.add(off0),
7046                *bytes_ptr.add(off0 + 1),
7047            ]));
7048            let s1 = f16_to_f32(u16::from_le_bytes([
7049                *bytes_ptr.add(off1),
7050                *bytes_ptr.add(off1 + 1),
7051            ]));
7052            let (u0_0, u1_0, u2_0, u3_0) = q1t_unpack_reg_u64s(bytes_ptr.add(off0 + 2));
7053            let (u0_1, u1_1, u2_1, u3_1) = q1t_unpack_reg_u64s(bytes_ptr.add(off1 + 2));
7054            let w0_0 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u0_0), vcreate_u64(u1_0)));
7055            let w1_0 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u2_0), vcreate_u64(u3_0)));
7056            let w0_1 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u0_1), vcreate_u64(u1_1)));
7057            let w1_1 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u2_1), vcreate_u64(u3_1)));
7058            for k in 0..2 {
7059                let d0 = sdot2!(w0_0, w1_0, xp[k].add(gi * GROUP_SIZE));
7060                let d1 = sdot2!(w0_1, w1_1, xp[k].add((gi + 1) * GROUP_SIZE));
7061                acc[k] += d0 as f32 * s0 + d1 as f32 * s1;
7062            }
7063            gi += 2;
7064        }
7065        if gi < gpr {
7066            let off = row_off + gi * TILE;
7067            let s = f16_to_f32(u16::from_le_bytes([
7068                *bytes_ptr.add(off),
7069                *bytes_ptr.add(off + 1),
7070            ]));
7071            let (u0, u1, u2, u3) = q1t_unpack_reg_u64s(bytes_ptr.add(off + 2));
7072            let w0 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u0), vcreate_u64(u1)));
7073            let w1 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u2), vcreate_u64(u3)));
7074            for k in 0..2 {
7075                let d = sdot2!(w0, w1, xp[k].add(gi * GROUP_SIZE));
7076                acc[k] += d as f32 * s;
7077            }
7078        }
7079        acc
7080    }
7081}
7082
7083/// Fused Q1T pair matvec: ONE pass over the rows serves both
7084/// activation streams — on ARM the ternary register unpack happens
7085/// once per tile pair (`q1t_dot_row_sdot2`); elsewhere the second dot
7086/// rides the row's L1-warm tile bytes. Per stream the math matches
7087/// `q1t_matvec` exactly.
7088fn q1t_matvec2(
7089    bytes: &[u8],
7090    x1: &[f32],
7091    x2: &[f32],
7092    rows: usize,
7093    cols: usize,
7094    o1: &mut [f32],
7095    o2: &mut [f32],
7096    pool: Option<&Pool>,
7097) {
7098    debug_assert_eq!(o1.len(), rows);
7099    debug_assert_eq!(o2.len(), rows);
7100    const TILE: usize = cortiq_core::quant::Q1T_TILE;
7101    let gpr = cols / GROUP_SIZE;
7102    let (rp_off, ent_off, has_ov) = q1t_overlay(bytes, rows * gpr * TILE, rows);
7103    let out1 = SendMut(o1.as_mut_ptr());
7104    let out2 = SendMut(o2.as_mut_ptr());
7105    if a8w8_enabled() {
7106        let a1 = split_act(x1);
7107        let a2 = split_act(x2);
7108        let (a1, a2) = (&a1, &a2);
7109        let run = move |start: usize, end: usize| {
7110            for r in start..end {
7111                #[cfg(target_arch = "aarch64")]
7112                // a8w8 on aarch64 ⇔ sdot_enabled(), so the kernel's
7113                // target features are present.
7114                let ds = unsafe { q1t_dot_row_sdot2(bytes, r, gpr, &a1.xq, &a2.xq) };
7115                #[cfg(not(target_arch = "aarch64"))]
7116                let ds = [
7117                    q1t_dot_row_i8(bytes, r, gpr, &a1.xq),
7118                    q1t_dot_row_i8(bytes, r, gpr, &a2.xq),
7119                ];
7120                let mut acc1 = ds[0] * a1.sx;
7121                for &(j, xv) in &a1.outliers {
7122                    acc1 += q1t_base_weight(bytes, r, gpr, j) * xv;
7123                }
7124                acc1 += q1t_row_outlier_correction(bytes, r, rp_off, ent_off, has_ov, x1);
7125                let mut acc2 = ds[1] * a2.sx;
7126                for &(j, xv) in &a2.outliers {
7127                    acc2 += q1t_base_weight(bytes, r, gpr, j) * xv;
7128                }
7129                acc2 += q1t_row_outlier_correction(bytes, r, rp_off, ent_off, has_ov, x2);
7130                // SAFETY: disjoint row ranges per worker.
7131                unsafe {
7132                    *out1.at(r) = acc1;
7133                    *out2.at(r) = acc2;
7134                }
7135            }
7136        };
7137        dispatch_rows(pool, rows, &run);
7138        return;
7139    }
7140    let run = move |start: usize, end: usize| {
7141        // Exact path (CMF_SDOT=0): unpack the sign LUT once per group,
7142        // dot both streams — same op order per stream as `q1t_matvec`.
7143        let mut sg = [0f32; GROUP_SIZE];
7144        for r in start..end {
7145            let mut acc1 = 0f32;
7146            let mut acc2 = 0f32;
7147            for g in 0..gpr {
7148                let off = (r * gpr + g) * TILE;
7149                let s = f16_to_f32(u16::from_le_bytes([bytes[off], bytes[off + 1]]));
7150                let codes = &bytes[off + 2..off + TILE];
7151                for bi in 0..6 {
7152                    sg[bi * 5..bi * 5 + 5].copy_from_slice(&SIGN5[codes[bi] as usize]);
7153                }
7154                let lut = &SIGN5[codes[6] as usize];
7155                sg[30] = lut[0];
7156                sg[31] = lut[1];
7157                let xg1 = &x1[g * GROUP_SIZE..g * GROUP_SIZE + GROUP_SIZE];
7158                let xg2 = &x2[g * GROUP_SIZE..g * GROUP_SIZE + GROUP_SIZE];
7159                let mut gsum1 = 0f32;
7160                for k in 0..GROUP_SIZE {
7161                    gsum1 += sg[k] * xg1[k];
7162                }
7163                acc1 += s * gsum1;
7164                let mut gsum2 = 0f32;
7165                for k in 0..GROUP_SIZE {
7166                    gsum2 += sg[k] * xg2[k];
7167                }
7168                acc2 += s * gsum2;
7169            }
7170            acc1 += q1t_row_outlier_correction(bytes, r, rp_off, ent_off, has_ov, x1);
7171            acc2 += q1t_row_outlier_correction(bytes, r, rp_off, ent_off, has_ov, x2);
7172            // SAFETY: disjoint row ranges per worker.
7173            unsafe {
7174                *out1.at(r) = acc1;
7175                *out2.at(r) = acc2;
7176            }
7177        }
7178    };
7179    dispatch_rows(pool, rows, &run);
7180}
7181
7182/// Ternary (q1t) matmat (prefill) — dequant each row once, dot the whole
7183/// batch against it (amortizes the per-row decode).
7184fn q1t_matmat(
7185    bytes: &[u8],
7186    xs: &[f32],
7187    b: usize,
7188    rows: usize,
7189    cols: usize,
7190    out: &mut [f32],
7191    pool: Option<&Pool>,
7192) {
7193    debug_assert_eq!(out.len(), b * rows);
7194    const TILE: usize = cortiq_core::quant::Q1T_TILE;
7195    let gpr = cols / GROUP_SIZE;
7196    let (rp_off, ent_off, has_ov) = q1t_overlay(bytes, rows * gpr * TILE, rows);
7197    let out_addr = SendMut(out.as_mut_ptr());
7198    // int8 prefill (ARM SDOT / x86 AVX2): quantize the B inputs once, unpack
7199    // each weight row's signs to i8 ONCE, then int8-dot against every input —
7200    // the row sign-decode amortizes over the whole batch. CMF_SDOT=0 → f32.
7201    if a8w8_enabled() {
7202        let acts: Vec<SplitAct> = (0..b)
7203            .map(|bi| split_act(&xs[bi * cols..(bi + 1) * cols]))
7204            .collect();
7205        let acts = &acts;
7206        let run = move |start: usize, end: usize| {
7207            let mut sg = vec![0i8; cols + 8]; // row signs, i8 (+8 unpack slack)
7208            let mut sc = vec![0f32; gpr]; // per-group scales
7209            let mut accs = vec![0f32; b]; // per-batch accumulators, reused per row
7210            for r in start..end {
7211                for g in 0..gpr {
7212                    let off = (r * gpr + g) * TILE;
7213                    sc[g] = f16_to_f32(u16::from_le_bytes([bytes[off], bytes[off + 1]]));
7214                    q1t_unpack_group_i8(
7215                        bytes.as_ptr().wrapping_add(off + 2),
7216                        &mut sg[g * GROUP_SIZE..],
7217                    );
7218                }
7219                for bi in 0..b {
7220                    let act = &acts[bi];
7221                    let mut isum = 0f32;
7222                    for g in 0..gpr {
7223                        let d = q1t_i8dot32(
7224                            sg.as_ptr().wrapping_add(g * GROUP_SIZE),
7225                            act.xq.as_ptr().wrapping_add(g * GROUP_SIZE),
7226                        );
7227                        isum += d as f32 * sc[g];
7228                    }
7229                    let mut acc = isum * act.sx;
7230                    for &(j, xv) in &act.outliers {
7231                        acc += q1t_base_weight(bytes, r, gpr, j) * xv;
7232                    }
7233                    accs[bi] = acc;
7234                }
7235                // Overlay ONCE per row for the whole batch: read each (col, val)
7236                // from mmap a single time (was b× — the re-read dominated prefill)
7237                // and fan it out over the batch via the cached inputs.
7238                if has_ov {
7239                    let (c0, c1) = (
7240                        q1t_rowptr(bytes, rp_off, r),
7241                        q1t_rowptr(bytes, rp_off, r + 1),
7242                    );
7243                    for p in c0..c1 {
7244                        let e = ent_off + p * 4;
7245                        let col = u16::from_le_bytes([bytes[e], bytes[e + 1]]) as usize;
7246                        let val = f16_to_f32(u16::from_le_bytes([bytes[e + 2], bytes[e + 3]]));
7247                        for bi in 0..b {
7248                            accs[bi] += val * xs[bi * cols + col];
7249                        }
7250                    }
7251                }
7252                for bi in 0..b {
7253                    unsafe { *out_addr.at(bi * rows + r) = accs[bi] };
7254                }
7255            }
7256        };
7257        dispatch_rows(pool, rows, &run);
7258        return;
7259    }
7260    let run = move |start: usize, end: usize| {
7261        let mut buf = vec![0f32; cols];
7262        for r in start..end {
7263            q1t_dequant_row(bytes, r, gpr, rp_off, ent_off, has_ov, &mut buf);
7264            for bi in 0..b {
7265                let xr = &xs[bi * cols..(bi + 1) * cols];
7266                let mut acc = 0f32;
7267                for j in 0..cols {
7268                    acc += buf[j] * xr[j];
7269                }
7270                unsafe { *out_addr.at(bi * rows + r) = acc };
7271            }
7272        }
7273    };
7274    dispatch_rows(pool, rows, &run);
7275}
7276
7277fn q1_matvec(
7278    bytes: &[u8],
7279    x: &[f32],
7280    rows: usize,
7281    cols: usize,
7282    out: &mut [f32],
7283    pool: Option<&Pool>,
7284) {
7285    debug_assert_eq!(out.len(), rows);
7286    let gpr = cols / GROUP_SIZE;
7287    let out_addr = SendMut(out.as_mut_ptr());
7288    if a8w8_enabled() {
7289        let act = split_act(x);
7290        let gsum = q1_group_sums(&act.xq, gpr);
7291        let (act, gsum) = (&act, &gsum);
7292        let run = move |start: usize, end: usize| {
7293            q1_range_a8w8(bytes, gpr, act, gsum, out_addr, start, end)
7294        };
7295        dispatch_rows(pool, rows, &run);
7296        return;
7297    }
7298    let run = move |start: usize, end: usize| q1_range_f32(bytes, gpr, x, out_addr, start, end);
7299    dispatch_rows(pool, rows, &run);
7300}
7301
7302/// Fused two-input q1 matvec (weights read once per pair).
7303#[allow(clippy::too_many_arguments)]
7304fn q1_matvec2(
7305    bytes: &[u8],
7306    x1: &[f32],
7307    x2: &[f32],
7308    rows: usize,
7309    cols: usize,
7310    o1: &mut [f32],
7311    o2: &mut [f32],
7312    pool: Option<&Pool>,
7313) {
7314    let gpr = cols / GROUP_SIZE;
7315    let p1 = SendMut(o1.as_mut_ptr());
7316    let p2 = SendMut(o2.as_mut_ptr());
7317    if a8w8_enabled() {
7318        let a1 = split_act(x1);
7319        let a2 = split_act(x2);
7320        let g1 = q1_group_sums(&a1.xq, gpr);
7321        let g2 = q1_group_sums(&a2.xq, gpr);
7322        let (a1, a2, g1, g2) = (&a1, &a2, &g1, &g2);
7323        let run = move |start: usize, end: usize| {
7324            for r in start..end {
7325                let mut v1 = dot_q1_row_i8(bytes, r, gpr, &a1.xq, g1) * a1.sx;
7326                let mut v2 = dot_q1_row_i8(bytes, r, gpr, &a2.xq, g2) * a2.sx;
7327                for &(j, xv) in &a1.outliers {
7328                    let (w, s) = q1_outlier(bytes, r, gpr, j);
7329                    v1 += w * s * xv;
7330                }
7331                for &(j, xv) in &a2.outliers {
7332                    let (w, s) = q1_outlier(bytes, r, gpr, j);
7333                    v2 += w * s * xv;
7334                }
7335                // SAFETY: disjoint row ranges per worker.
7336                unsafe {
7337                    *p1.at(r) = v1;
7338                    *p2.at(r) = v2;
7339                }
7340            }
7341        };
7342        dispatch_rows(pool, rows, &run);
7343        return;
7344    }
7345    let run = move |start: usize, end: usize| {
7346        for r in start..end {
7347            // SAFETY: disjoint row ranges per worker.
7348            unsafe {
7349                *p1.at(r) = q1_row_exact(bytes, r, gpr, x1);
7350                *p2.at(r) = q1_row_exact(bytes, r, gpr, x2);
7351            }
7352        }
7353    };
7354    dispatch_rows(pool, rows, &run);
7355}
7356
7357/// Batched q1 matmat: each row's tiles stream once per microbatch.
7358#[allow(clippy::too_many_arguments)]
7359fn q1_matmat(
7360    bytes: &[u8],
7361    xs_all: &[f32],
7362    b: usize,
7363    rows: usize,
7364    cols: usize,
7365    out: &mut [f32],
7366    pool: Option<&Pool>,
7367) {
7368    debug_assert_eq!(out.len(), b * rows);
7369    let gpr = cols / GROUP_SIZE;
7370    let out_addr = SendMut(out.as_mut_ptr());
7371    if a8w8_enabled() {
7372        let acts: Vec<(SplitAct, Vec<i32>)> = (0..b)
7373            .map(|bi| {
7374                let act = split_act(&xs_all[bi * cols..(bi + 1) * cols]);
7375                let gsum = q1_group_sums(&act.xq, gpr);
7376                (act, gsum)
7377            })
7378            .collect();
7379        let acts = &acts;
7380        #[cfg(target_arch = "x86_64")]
7381        let blocked_ok = avx2_enabled() && blocked_enabled();
7382        #[cfg(target_arch = "aarch64")]
7383        let blocked_ok = sdot_enabled() && blocked_enabled();
7384        let run = move |start: usize, end: usize| {
7385            for r in start..end {
7386                let mut bi = 0usize;
7387                // Blocked 1×4: the unpacked bit mask serves four
7388                // activation streams per group.
7389                #[cfg(target_arch = "aarch64")]
7390                if blocked_ok {
7391                    while bi + 4 <= acts.len() {
7392                        let xs = [
7393                            acts[bi].0.xq.as_slice(),
7394                            acts[bi + 1].0.xq.as_slice(),
7395                            acts[bi + 2].0.xq.as_slice(),
7396                            acts[bi + 3].0.xq.as_slice(),
7397                        ];
7398                        let gs = [
7399                            acts[bi].1.as_slice(),
7400                            acts[bi + 1].1.as_slice(),
7401                            acts[bi + 2].1.as_slice(),
7402                            acts[bi + 3].1.as_slice(),
7403                        ];
7404                        let d = unsafe { dot_q1_row_1x4_sdot(bytes, r, gpr, xs, gs) };
7405                        for k in 0..4 {
7406                            let (act, _) = &acts[bi + k];
7407                            let mut acc = d[k] * act.sx;
7408                            for &(j, xv) in &act.outliers {
7409                                let (w, sc) = q1_outlier(bytes, r, gpr, j);
7410                                acc += w * sc * xv;
7411                            }
7412                            // SAFETY: disjoint (bi, r) cells per worker.
7413                            unsafe { *out_addr.at((bi + k) * rows + r) = acc };
7414                        }
7415                        bi += 4;
7416                    }
7417                }
7418                #[cfg(target_arch = "x86_64")]
7419                if blocked_ok {
7420                    while bi + 4 <= acts.len() {
7421                        let xs = [
7422                            acts[bi].0.xq.as_slice(),
7423                            acts[bi + 1].0.xq.as_slice(),
7424                            acts[bi + 2].0.xq.as_slice(),
7425                            acts[bi + 3].0.xq.as_slice(),
7426                        ];
7427                        let gs = [
7428                            acts[bi].1.as_slice(),
7429                            acts[bi + 1].1.as_slice(),
7430                            acts[bi + 2].1.as_slice(),
7431                            acts[bi + 3].1.as_slice(),
7432                        ];
7433                        let d = unsafe {
7434                            if vnni_tiles_enabled() {
7435                                dot_q1_row_1x4_vnni(bytes, r, gpr, xs, gs)
7436                            } else {
7437                                dot_q1_row_1x4_avx2(bytes, r, gpr, xs, gs)
7438                            }
7439                        };
7440                        for k in 0..4 {
7441                            let (act, _) = &acts[bi + k];
7442                            let mut acc = d[k] * act.sx;
7443                            for &(j, xv) in &act.outliers {
7444                                let (w, sc) = q1_outlier(bytes, r, gpr, j);
7445                                acc += w * sc * xv;
7446                            }
7447                            // SAFETY: disjoint (bi, r) cells per worker.
7448                            unsafe { *out_addr.at((bi + k) * rows + r) = acc };
7449                        }
7450                        bi += 4;
7451                    }
7452                }
7453                while bi < acts.len() {
7454                    let (act, gsum) = &acts[bi];
7455                    let mut acc = dot_q1_row_i8(bytes, r, gpr, &act.xq, gsum) * act.sx;
7456                    for &(j, xv) in &act.outliers {
7457                        let (w, s) = q1_outlier(bytes, r, gpr, j);
7458                        acc += w * s * xv;
7459                    }
7460                    // SAFETY: disjoint (bi, r) cells per worker range.
7461                    unsafe { *out_addr.at(bi * rows + r) = acc };
7462                    bi += 1;
7463                }
7464            }
7465        };
7466        dispatch_rows(pool, rows, &run);
7467        return;
7468    }
7469    let run = move |start: usize, end: usize| {
7470        for r in start..end {
7471            for bi in 0..b {
7472                let x = &xs_all[bi * cols..(bi + 1) * cols];
7473                // SAFETY: disjoint (bi, r) cells per worker range.
7474                unsafe { *out_addr.at(bi * rows + r) = q1_row_exact(bytes, r, gpr, x) };
7475            }
7476        }
7477    };
7478    dispatch_rows(pool, rows, &run);
7479}
7480
7481/// Fused q4_block matvec straight from the mapped bytes. SDOT path when
7482/// dotprod is available (port of vmfcore `dot_q4_block_sdot`, measured
7483/// +23% on q4 decode): nibbles → centered i8, int8×int8 `sdot` per
7484/// 32-group, exact outlier correction — the same A8W8 contract as q8.
7485/// `CMF_SDOT=0` keeps the exact scalar path.
7486fn q4matvec(
7487    bytes: &[u8],
7488    x: &[f32],
7489    rows: usize,
7490    cols: usize,
7491    out: &mut [f32],
7492    pool: Option<&Pool>,
7493) {
7494    debug_assert_eq!(out.len(), rows);
7495    let (packed, scales) = q4_split(bytes, rows, cols);
7496    let gpr = cols / GROUP_SIZE;
7497    let out_addr = SendMut(out.as_mut_ptr());
7498
7499    if a8w8_enabled() {
7500        let act = split_act(x);
7501        let run = move |start: usize, end: usize| {
7502            q4_range_a8w8(packed, scales, gpr, cols, &act, out_addr, start, end)
7503        };
7504        dispatch_rows(pool, rows, &run);
7505        return;
7506    }
7507
7508    let run =
7509        move |start: usize, end: usize| q4_range_f32(packed, scales, gpr, x, out_addr, start, end);
7510    dispatch_rows(pool, rows, &run);
7511}
7512
7513/// One q4 row via the A8W8 int8 path — SDOT on ARM, AVX2 maddubs on
7514/// x86 (scalar fallback is unreachable: callers gate on a8w8_enabled).
7515#[inline]
7516#[allow(unreachable_code)]
7517/// One UNPACKED q4 row (centered i8 in `buf`) against four activation
7518/// streams: the 32-byte weight chunk and its abs() load once per group,
7519/// the per-group f16 scale decodes once — four maddubs+reduce chains
7520/// instead of four full (load, abs, dot) rounds.
7521#[cfg(target_arch = "x86_64")]
7522#[target_feature(enable = "avx2")]
7523unsafe fn dot_q4b_row_1x4_avx2(
7524    buf: &[u8],
7525    scales: &[u8],
7526    g0: usize,
7527    gpr: usize,
7528    xs: [&[i8]; 4],
7529) -> [f32; 4] {
7530    // SAFETY: callers uphold buffer contracts (buf.len() == gpr·32).
7531    unsafe {
7532        use core::arch::x86_64::*;
7533        let ones = _mm256_set1_epi16(1);
7534        let mut acc = [0f32; 4];
7535        for gi in 0..gpr {
7536            let s = f16_to_f32(u16::from_le_bytes([
7537                scales[(g0 + gi) * 2],
7538                scales[(g0 + gi) * 2 + 1],
7539            ]));
7540            let w = _mm256_loadu_si256(buf.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
7541            let aw = _mm256_abs_epi8(w);
7542            for (k, xq) in xs.iter().enumerate() {
7543                let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
7544                let p16 = _mm256_maddubs_epi16(aw, _mm256_sign_epi8(x, w));
7545                let d = _mm256_madd_epi16(p16, ones);
7546                let hi128 = _mm256_extracti128_si256::<1>(d);
7547                let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
7548                let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
7549                let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
7550                acc[k] += _mm_cvtsi128_si32(s32) as f32 * s;
7551            }
7552        }
7553        acc
7554    }
7555}
7556
7557/// VNNI twin of `dot_q4b_row_1x4_avx2` (see `dpbusd_hsum`).
7558#[cfg(target_arch = "x86_64")]
7559#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
7560unsafe fn dot_q4b_row_1x4_vnni(
7561    buf: &[u8],
7562    scales: &[u8],
7563    g0: usize,
7564    gpr: usize,
7565    xs: [&[i8]; 4],
7566) -> [f32; 4] {
7567    // SAFETY: callers uphold buffer contracts (buf.len() == gpr·32).
7568    unsafe {
7569        use core::arch::x86_64::*;
7570        let mut acc = [0f32; 4];
7571        for gi in 0..gpr {
7572            let s = f16_to_f32(u16::from_le_bytes([
7573                scales[(g0 + gi) * 2],
7574                scales[(g0 + gi) * 2 + 1],
7575            ]));
7576            let w = _mm256_loadu_si256(buf.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
7577            let aw = _mm256_abs_epi8(w);
7578            for (k, xq) in xs.iter().enumerate() {
7579                let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
7580                let d = dpbusd_hsum(aw, _mm256_sign_epi8(x, w));
7581                acc[k] += d as f32 * s;
7582            }
7583        }
7584        acc
7585    }
7586}
7587
7588/// The vbit flavor of the blocked 1×4: the per-activation A8W8 scale
7589/// folds in PER GROUP as `(d·sx)·s` — bit-matching the single-matvec
7590/// accumulation order (the q4_block flavor applies sx once at the end,
7591/// matching ITS single path; the two conventions are historical and
7592/// each blocked leg must mirror its own).
7593#[cfg(target_arch = "x86_64")]
7594#[target_feature(enable = "avx2")]
7595unsafe fn dot_q4b_row_1x4_sx_avx2(
7596    buf: &[u8],
7597    scales: &[u8],
7598    g0: usize,
7599    gpr: usize,
7600    xs: [&[i8]; 4],
7601    sxs: [f32; 4],
7602) -> [f32; 4] {
7603    // SAFETY: callers uphold buffer contracts (buf.len() == gpr·32).
7604    unsafe {
7605        use core::arch::x86_64::*;
7606        let ones = _mm256_set1_epi16(1);
7607        let mut acc = [0f32; 4];
7608        for gi in 0..gpr {
7609            let s = f16_to_f32(u16::from_le_bytes([
7610                scales[(g0 + gi) * 2],
7611                scales[(g0 + gi) * 2 + 1],
7612            ]));
7613            let w = _mm256_loadu_si256(buf.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
7614            let aw = _mm256_abs_epi8(w);
7615            for (k, xq) in xs.iter().enumerate() {
7616                let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
7617                let p16 = _mm256_maddubs_epi16(aw, _mm256_sign_epi8(x, w));
7618                let d = _mm256_madd_epi16(p16, ones);
7619                let hi128 = _mm256_extracti128_si256::<1>(d);
7620                let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
7621                let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
7622                let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
7623                acc[k] += (_mm_cvtsi128_si32(s32) as f32 * sxs[k]) * s;
7624            }
7625        }
7626        acc
7627    }
7628}
7629
7630/// VNNI twin of `dot_q4b_row_1x4_sx_avx2` (see `dpbusd_hsum`; the
7631/// per-group `(d·sx)·s` fold mirrors the vbit single path).
7632#[cfg(target_arch = "x86_64")]
7633#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
7634unsafe fn dot_q4b_row_1x4_sx_vnni(
7635    buf: &[u8],
7636    scales: &[u8],
7637    g0: usize,
7638    gpr: usize,
7639    xs: [&[i8]; 4],
7640    sxs: [f32; 4],
7641) -> [f32; 4] {
7642    // SAFETY: callers uphold buffer contracts (buf.len() == gpr·32).
7643    unsafe {
7644        use core::arch::x86_64::*;
7645        let mut acc = [0f32; 4];
7646        for gi in 0..gpr {
7647            let s = f16_to_f32(u16::from_le_bytes([
7648                scales[(g0 + gi) * 2],
7649                scales[(g0 + gi) * 2 + 1],
7650            ]));
7651            let w = _mm256_loadu_si256(buf.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
7652            let aw = _mm256_abs_epi8(w);
7653            for (k, xq) in xs.iter().enumerate() {
7654                let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
7655                let d = dpbusd_hsum(aw, _mm256_sign_epi8(x, w));
7656                acc[k] += (d as f32 * sxs[k]) * s;
7657            }
7658        }
7659        acc
7660    }
7661}
7662
7663#[allow(unreachable_code)]
7664fn dot_q4_row_i8(packed: &[u8], scales: &[u8], g0: usize, gpr: usize, xq: &[i8]) -> f32 {
7665    #[cfg(target_arch = "aarch64")]
7666    unsafe {
7667        return dot_q4_row_sdot(packed, scales, g0, gpr, xq);
7668    }
7669    #[cfg(target_arch = "x86_64")]
7670    unsafe {
7671        return dot_q4_row_avx2(packed, scales, g0, gpr, xq);
7672    }
7673    let mut acc = 0f32;
7674    for gi in 0..gpr {
7675        let g = g0 + gi;
7676        let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
7677        let mut d = 0i32;
7678        for (k, &b) in packed[g * 16..(g + 1) * 16].iter().enumerate() {
7679            d += ((b & 0x0F) as i32 - 8) * xq[gi * GROUP_SIZE + k * 2] as i32
7680                + (((b >> 4) & 0x0F) as i32 - 8) * xq[gi * GROUP_SIZE + k * 2 + 1] as i32;
7681        }
7682        acc += d as f32 * s;
7683    }
7684    acc
7685}
7686
7687/// Two-activation q4 row via the A8W8 int8 path (see `dot_q4_row_i8`).
7688#[inline]
7689#[allow(unreachable_code)]
7690fn dot_q4_row_i8_2(
7691    packed: &[u8],
7692    scales: &[u8],
7693    g0: usize,
7694    gpr: usize,
7695    xq1: &[i8],
7696    xq2: &[i8],
7697) -> (f32, f32) {
7698    #[cfg(target_arch = "aarch64")]
7699    unsafe {
7700        return dot_q4_row_sdot2(packed, scales, g0, gpr, xq1, xq2);
7701    }
7702    #[cfg(target_arch = "x86_64")]
7703    unsafe {
7704        return dot_q4_row_avx2_2(packed, scales, g0, gpr, xq1, xq2);
7705    }
7706    (
7707        dot_q4_row_i8(packed, scales, g0, gpr, xq1),
7708        dot_q4_row_i8(packed, scales, g0, gpr, xq2),
7709    )
7710}
7711
7712/// One q4 row range via SDOT (kernel body of `q4matvec`, extracted so
7713/// multi-matrix jobs can drive it for several tensors in one dispatch).
7714#[allow(clippy::too_many_arguments)]
7715fn q4_range_a8w8(
7716    packed: &[u8],
7717    scales: &[u8],
7718    gpr: usize,
7719    cols: usize,
7720    act: &SplitAct,
7721    out: SendMut,
7722    start: usize,
7723    end: usize,
7724) {
7725    for r in start..end {
7726        let mut acc = dot_q4_row_i8(packed, scales, r * gpr, gpr, &act.xq) * act.sx;
7727        // xq is zeroed at outlier slots — add the exact terms.
7728        for &(j, xv) in &act.outliers {
7729            let flat = r * cols + j;
7730            let byte = packed[flat / 2];
7731            let nib = if flat & 1 == 0 {
7732                byte & 0x0F
7733            } else {
7734                byte >> 4
7735            };
7736            let s = f16_to_f32(u16::from_le_bytes([
7737                scales[(flat / GROUP_SIZE) * 2],
7738                scales[(flat / GROUP_SIZE) * 2 + 1],
7739            ]));
7740            acc += ((nib as i32 - 8) as f32) * s * xv;
7741        }
7742        // SAFETY: disjoint row ranges per worker.
7743        unsafe { *out.at(r) = acc };
7744    }
7745}
7746
7747/// Two-input q4 row range via the A8W8 int8 path — kernel body of
7748/// `q4matvec2`, extracted for pair multi-matrix jobs.
7749#[allow(clippy::too_many_arguments)]
7750fn q4_range2_a8w8(
7751    packed: &[u8],
7752    scales: &[u8],
7753    gpr: usize,
7754    cols: usize,
7755    a1: &SplitAct,
7756    a2: &SplitAct,
7757    p1: SendMut,
7758    p2: SendMut,
7759    start: usize,
7760    end: usize,
7761) {
7762    for r in start..end {
7763        let (s1, s2) = dot_q4_row_i8_2(packed, scales, r * gpr, gpr, &a1.xq, &a2.xq);
7764        let mut acc1 = s1 * a1.sx;
7765        let mut acc2 = s2 * a2.sx;
7766        // xq is zeroed at outlier slots — add the exact terms.
7767        let fix = |outliers: &[(usize, f32)], acc: &mut f32| {
7768            for &(j, xv) in outliers {
7769                let flat = r * cols + j;
7770                let byte = packed[flat / 2];
7771                let nib = if flat & 1 == 0 {
7772                    byte & 0x0F
7773                } else {
7774                    byte >> 4
7775                };
7776                let s = f16_to_f32(u16::from_le_bytes([
7777                    scales[(flat / GROUP_SIZE) * 2],
7778                    scales[(flat / GROUP_SIZE) * 2 + 1],
7779                ]));
7780                *acc += ((nib as i32 - 8) as f32) * s * xv;
7781            }
7782        };
7783        fix(&a1.outliers, &mut acc1);
7784        fix(&a2.outliers, &mut acc2);
7785        // SAFETY: disjoint row ranges per worker.
7786        unsafe {
7787            *p1.at(r) = acc1;
7788            *p2.at(r) = acc2;
7789        }
7790    }
7791}
7792
7793/// Exact scalar q4 row range (same extraction, non-SDOT path).
7794fn q4_range_f32(
7795    packed: &[u8],
7796    scales: &[u8],
7797    gpr: usize,
7798    x: &[f32],
7799    out: SendMut,
7800    start: usize,
7801    end: usize,
7802) {
7803    for r in start..end {
7804        let mut acc = 0f32;
7805        for gi in 0..gpr {
7806            let g = r * gpr + gi;
7807            let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
7808            let pk = &packed[g * 16..(g + 1) * 16];
7809            let xg = &x[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
7810            let mut ga = 0f32;
7811            for (k, &b) in pk.iter().enumerate() {
7812                ga += ((b & 0x0F) as f32 - 8.0) * xg[k * 2]
7813                    + (((b >> 4) & 0x0F) as f32 - 8.0) * xg[k * 2 + 1];
7814            }
7815            acc += ga * s;
7816        }
7817        // SAFETY: disjoint row ranges per worker.
7818        unsafe { *out.at(r) = acc };
7819    }
7820}
7821
7822/// Fused two-input q4 matvec: nibbles are unpacked ONCE per group and
7823/// dotted against both activations (was: two full matvecs — double
7824/// weight traffic). Per-lane math matches `q4matvec` exactly.
7825#[allow(clippy::too_many_arguments)]
7826fn q4matvec2(
7827    bytes: &[u8],
7828    x1: &[f32],
7829    x2: &[f32],
7830    rows: usize,
7831    cols: usize,
7832    o1: &mut [f32],
7833    o2: &mut [f32],
7834    pool: Option<&Pool>,
7835) {
7836    debug_assert_eq!(o1.len(), rows);
7837    debug_assert_eq!(o2.len(), rows);
7838    let (packed, scales) = q4_split(bytes, rows, cols);
7839    let gpr = cols / GROUP_SIZE;
7840
7841    if a8w8_enabled() {
7842        let a1 = split_act(x1);
7843        let a2 = split_act(x2);
7844        let p1 = SendMut(o1.as_mut_ptr());
7845        let p2 = SendMut(o2.as_mut_ptr());
7846        let run = move |start: usize, end: usize| {
7847            q4_range2_a8w8(packed, scales, gpr, cols, &a1, &a2, p1, p2, start, end)
7848        };
7849        dispatch_rows(pool, rows, &run);
7850        return;
7851    }
7852
7853    let p1 = SendMut(o1.as_mut_ptr());
7854    let p2 = SendMut(o2.as_mut_ptr());
7855    let run = move |start: usize, end: usize| {
7856        q4_range2_f32(packed, scales, gpr, x1, x2, p1, p2, start, end)
7857    };
7858    dispatch_rows(pool, rows, &run);
7859}
7860
7861/// Two-input exact scalar q4 row range (same extraction).
7862#[allow(clippy::too_many_arguments)]
7863fn q4_range2_f32(
7864    packed: &[u8],
7865    scales: &[u8],
7866    gpr: usize,
7867    x1: &[f32],
7868    x2: &[f32],
7869    p1: SendMut,
7870    p2: SendMut,
7871    start: usize,
7872    end: usize,
7873) {
7874    for r in start..end {
7875        let (mut acc1, mut acc2) = (0f32, 0f32);
7876        for gi in 0..gpr {
7877            let g = r * gpr + gi;
7878            let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
7879            let pk = &packed[g * 16..(g + 1) * 16];
7880            let x1g = &x1[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
7881            let x2g = &x2[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
7882            let (mut g1, mut g2) = (0f32, 0f32);
7883            for (k, &b) in pk.iter().enumerate() {
7884                let wl = (b & 0x0F) as f32 - 8.0;
7885                let wh = ((b >> 4) & 0x0F) as f32 - 8.0;
7886                g1 += wl * x1g[k * 2] + wh * x1g[k * 2 + 1];
7887                g2 += wl * x2g[k * 2] + wh * x2g[k * 2 + 1];
7888            }
7889            acc1 += g1 * s;
7890            acc2 += g2 * s;
7891        }
7892        // SAFETY: disjoint row ranges per worker.
7893        unsafe {
7894            *p1.at(r) = acc1;
7895            *p2.at(r) = acc2;
7896        }
7897    }
7898}
7899
7900thread_local! {
7901    /// Per-worker decoded-row scratch for the batched q4/vbit kernels
7902    /// (centered i8 for SDOT, f32 for the exact/scalar paths).
7903    static ROW_I8: std::cell::RefCell<Vec<u8>> = const { std::cell::RefCell::new(Vec::new()) };
7904    static ROW_F32: std::cell::RefCell<Vec<f32>> = const { std::cell::RefCell::new(Vec::new()) };
7905}
7906
7907/// Batched q4 matmat: each weight row is unpacked from the mmap ONCE
7908/// and dotted against ALL b activations (prefill used to fall back to b
7909/// full matvecs — b× weight traffic and b× nibble decode). Per-position
7910/// math matches `q4matvec` exactly: same group order, same accumulation.
7911/// `out` is row-major [b, rows] like `qmatmat`.
7912#[allow(clippy::too_many_arguments)]
7913fn q4matmat(
7914    bytes: &[u8],
7915    xs_all: &[f32],
7916    b: usize,
7917    rows: usize,
7918    cols: usize,
7919    out: &mut [f32],
7920    pool: Option<&Pool>,
7921) {
7922    debug_assert_eq!(xs_all.len(), b * cols);
7923    debug_assert_eq!(out.len(), b * rows);
7924    let (packed, scales) = q4_split(bytes, rows, cols);
7925    let gpr = cols / GROUP_SIZE;
7926    let gscale = |g: usize| f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
7927
7928    if a8w8_enabled() {
7929        let acts: Vec<SplitAct> = (0..b)
7930            .map(|bi| split_act(&xs_all[bi * cols..(bi + 1) * cols]))
7931            .collect();
7932        let acts = &acts;
7933        let out_addr = SendMut(out.as_mut_ptr());
7934        let run = move |start: usize, end: usize| {
7935            ROW_I8.with(|rb| {
7936                let mut buf = rb.borrow_mut();
7937                buf.resize(cols, 0);
7938                for r in start..end {
7939                    // Unpack the row's nibbles to centered i8 once
7940                    // (element 2k = low nibble, 2k+1 = high — flat order,
7941                    // same as dot_q4_row_sdot's zip).
7942                    for gi in 0..gpr {
7943                        let g = r * gpr + gi;
7944                        for (k, &bt) in packed[g * 16..(g + 1) * 16].iter().enumerate() {
7945                            buf[gi * GROUP_SIZE + k * 2] = ((bt & 0x0F) as i32 - 8) as i8 as u8;
7946                            buf[gi * GROUP_SIZE + k * 2 + 1] =
7947                                (((bt >> 4) & 0x0F) as i32 - 8) as i8 as u8;
7948                        }
7949                    }
7950                    let mut bi = 0usize;
7951                    #[cfg(target_arch = "x86_64")]
7952                    if avx2_enabled() && blocked_enabled() {
7953                        while bi + 4 <= acts.len() {
7954                            let xs = [
7955                                acts[bi].xq.as_slice(),
7956                                acts[bi + 1].xq.as_slice(),
7957                                acts[bi + 2].xq.as_slice(),
7958                                acts[bi + 3].xq.as_slice(),
7959                            ];
7960                            let d = unsafe {
7961                                if vnni_tiles_enabled() {
7962                                    dot_q4b_row_1x4_vnni(&buf, scales, r * gpr, gpr, xs)
7963                                } else {
7964                                    dot_q4b_row_1x4_avx2(&buf, scales, r * gpr, gpr, xs)
7965                                }
7966                            };
7967                            for k in 0..4 {
7968                                let act = &acts[bi + k];
7969                                let mut acc = d[k] * act.sx;
7970                                for &(j, xv) in &act.outliers {
7971                                    acc += (buf[j] as i8) as f32
7972                                        * gscale((r * cols + j) / GROUP_SIZE)
7973                                        * xv;
7974                                }
7975                                // SAFETY: disjoint (bi, r) cells per worker.
7976                                unsafe { *out_addr.at((bi + k) * rows + r) = acc };
7977                            }
7978                            bi += 4;
7979                        }
7980                    }
7981                    while bi < acts.len() {
7982                        let act = &acts[bi];
7983                        let mut acc = 0f32;
7984                        for gi in 0..gpr {
7985                            let d = dot_i8_i8(
7986                                &buf[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE],
7987                                &act.xq[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE],
7988                            );
7989                            acc += d as f32 * gscale(r * gpr + gi);
7990                        }
7991                        acc *= act.sx;
7992                        // xq is zeroed at outlier slots — exact terms.
7993                        for &(j, xv) in &act.outliers {
7994                            acc += (buf[j] as i8) as f32 * gscale((r * cols + j) / GROUP_SIZE) * xv;
7995                        }
7996                        // SAFETY: disjoint (bi, r) cells per worker row range.
7997                        unsafe { *out_addr.at(bi * rows + r) = acc };
7998                        bi += 1;
7999                    }
8000                }
8001            })
8002        };
8003        dispatch_rows(pool, rows, &run);
8004        return;
8005    }
8006
8007    let out_addr = SendMut(out.as_mut_ptr());
8008    let run = move |start: usize, end: usize| {
8009        ROW_F32.with(|rb| {
8010            let mut buf = rb.borrow_mut();
8011            buf.resize(cols, 0.0);
8012            for r in start..end {
8013                // Decode raw (nib − 8) values once; scales stay per-group
8014                // so the accumulation order matches q4matvec bit-for-bit.
8015                for gi in 0..gpr {
8016                    let g = r * gpr + gi;
8017                    for (k, &bt) in packed[g * 16..(g + 1) * 16].iter().enumerate() {
8018                        buf[gi * GROUP_SIZE + k * 2] = (bt & 0x0F) as f32 - 8.0;
8019                        buf[gi * GROUP_SIZE + k * 2 + 1] = ((bt >> 4) & 0x0F) as f32 - 8.0;
8020                    }
8021                }
8022                for bi in 0..b {
8023                    let x = &xs_all[bi * cols..(bi + 1) * cols];
8024                    let mut acc = 0f32;
8025                    for gi in 0..gpr {
8026                        let mut ga = 0f32;
8027                        // Pairwise (lo + hi) addition, matching
8028                        // q4matvec's `ga += lo·x + hi·x` shape exactly —
8029                        // a flat one-per-element loop rounds differently
8030                        // and broke bit-parity on the scalar (x86) path.
8031                        for k in 0..GROUP_SIZE / 2 {
8032                            let e = gi * GROUP_SIZE + k * 2;
8033                            ga += buf[e] * x[e] + buf[e + 1] * x[e + 1];
8034                        }
8035                        acc += ga * gscale(r * gpr + gi);
8036                    }
8037                    // SAFETY: disjoint (bi, r) cells per worker row range.
8038                    unsafe { *out_addr.at(bi * rows + r) = acc };
8039                }
8040            }
8041        })
8042    };
8043    dispatch_rows(pool, rows, &run);
8044}
8045
8046/// Batched vbit matmat: each variable-bit row is decoded from the mmap
8047/// ONCE for the whole microbatch. Same per-position math as
8048/// `vbitmatvec` (SDOT A8W8 with exact outliers / exact f32 for b=8 rows
8049/// and the scalar path).
8050#[allow(clippy::too_many_arguments)]
8051fn vbitmatmat(
8052    bytes: &[u8],
8053    offsets: &[usize],
8054    xs_all: &[f32],
8055    b: usize,
8056    rows: usize,
8057    cols: usize,
8058    out: &mut [f32],
8059    pool: Option<&Pool>,
8060) {
8061    debug_assert_eq!(xs_all.len(), b * cols);
8062    debug_assert_eq!(out.len(), b * rows);
8063    debug_assert_eq!(offsets.len(), rows + 1);
8064    let ng = cols / GROUP_SIZE;
8065    let bits = &bytes[..rows];
8066    let sc_off = rows;
8067    let gscale = |r: usize, g: usize| {
8068        let so = (r * ng + g) * 2;
8069        f16_to_f32(u16::from_le_bytes([
8070            bytes[sc_off + so],
8071            bytes[sc_off + so + 1],
8072        ]))
8073    };
8074
8075    // Decode row r's raw (u − L) values into `dst` (f32, unscaled).
8076    let decode_f32 = |r: usize, dst: &mut [f32]| {
8077        let bw = bits[r] as usize;
8078        let l = ((1i32 << (bw - 1)) - 1) as f32;
8079        let data = &bytes[offsets[r]..offsets[r + 1]];
8080        let (mut acc, mut nbits, mut idx) = (0u64, 0usize, 0usize);
8081        for d in dst.iter_mut() {
8082            while nbits < bw {
8083                acc = (acc << 8) | data[idx] as u64;
8084                idx += 1;
8085                nbits += 8;
8086            }
8087            let u = ((acc >> (nbits - bw)) & ((1u64 << bw) - 1)) as f32;
8088            nbits -= bw;
8089            *d = u - l;
8090        }
8091    };
8092
8093    if a8w8_enabled() {
8094        let acts: Vec<SplitAct> = (0..b)
8095            .map(|bi| split_act(&xs_all[bi * cols..(bi + 1) * cols]))
8096            .collect();
8097        let acts = &acts;
8098        let out_addr = SendMut(out.as_mut_ptr());
8099        let run = move |start: usize, end: usize| {
8100            for r in start..end {
8101                let bw = bits[r] as usize;
8102                if bw == 8 {
8103                    // u−L reaches 128 → no i8 path; decode once, exact
8104                    // f32 dots for every position (same as vbitmatvec).
8105                    ROW_F32.with(|rb| {
8106                        let mut buf = rb.borrow_mut();
8107                        buf.resize(cols, 0.0);
8108                        decode_f32(r, &mut buf);
8109                        for bi in 0..b {
8110                            let x = &xs_all[bi * cols..(bi + 1) * cols];
8111                            let mut dot = 0f32;
8112                            for g in 0..ng {
8113                                let mut gd = 0f32;
8114                                for k in 0..GROUP_SIZE {
8115                                    gd += buf[g * GROUP_SIZE + k] * x[g * GROUP_SIZE + k];
8116                                }
8117                                dot += gd * gscale(r, g);
8118                            }
8119                            // SAFETY: disjoint (bi, r) cells per worker range.
8120                            unsafe { *out_addr.at(bi * rows + r) = dot };
8121                        }
8122                    });
8123                    continue;
8124                }
8125                let l = (1i32 << (bw - 1)) - 1;
8126                let data = &bytes[offsets[r]..offsets[r + 1]];
8127                ROW_I8.with(|rb| {
8128                    let mut buf = rb.borrow_mut();
8129                    buf.resize(cols, 0);
8130                    #[inline(always)]
8131                    fn fill<const B: usize>(data: &[u8], l: i32, buf: &mut [u8]) {
8132                        for (blk, chunk) in buf.chunks_exact_mut(8).enumerate() {
8133                            let u = unpack8::<B>(&data[blk * B..]);
8134                            for k in 0..8 {
8135                                chunk[k] = (u[k] - l) as i8 as u8;
8136                            }
8137                        }
8138                    }
8139                    match bw {
8140                        3 => fill::<3>(data, l, &mut buf),
8141                        4 => vbit_fill4(data, &mut buf),
8142                        5 => fill::<5>(data, l, &mut buf),
8143                        6 => fill::<6>(data, l, &mut buf),
8144                        _ => unreachable!("vbit bit-width {bw} (validated at load)"),
8145                    }
8146                    let mut bi = 0usize;
8147                    // The vbit scale table shares q4_block's layout
8148                    // (contiguous f16 per (row·ng + g)), so the same
8149                    // blocked 1×4 kernel serves the decoded row.
8150                    #[cfg(target_arch = "x86_64")]
8151                    if avx2_enabled() && blocked_enabled() {
8152                        while bi + 4 <= acts.len() {
8153                            let xs = [
8154                                acts[bi].xq.as_slice(),
8155                                acts[bi + 1].xq.as_slice(),
8156                                acts[bi + 2].xq.as_slice(),
8157                                acts[bi + 3].xq.as_slice(),
8158                            ];
8159                            let sxs = [
8160                                acts[bi].sx,
8161                                acts[bi + 1].sx,
8162                                acts[bi + 2].sx,
8163                                acts[bi + 3].sx,
8164                            ];
8165                            let d = unsafe {
8166                                if vnni_tiles_enabled() {
8167                                    dot_q4b_row_1x4_sx_vnni(
8168                                        &buf,
8169                                        &bytes[sc_off..],
8170                                        r * ng,
8171                                        ng,
8172                                        xs,
8173                                        sxs,
8174                                    )
8175                                } else {
8176                                    dot_q4b_row_1x4_sx_avx2(
8177                                        &buf,
8178                                        &bytes[sc_off..],
8179                                        r * ng,
8180                                        ng,
8181                                        xs,
8182                                        sxs,
8183                                    )
8184                                }
8185                            };
8186                            for k in 0..4 {
8187                                let act = &acts[bi + k];
8188                                let mut dot = d[k];
8189                                for &(j, xv) in &act.outliers {
8190                                    dot += (buf[j] as i8) as f32 * gscale(r, j / GROUP_SIZE) * xv;
8191                                }
8192                                // SAFETY: disjoint (bi, r) cells per worker.
8193                                unsafe { *out_addr.at((bi + k) * rows + r) = dot };
8194                            }
8195                            bi += 4;
8196                        }
8197                    }
8198                    while bi < acts.len() {
8199                        let act = &acts[bi];
8200                        let mut dot = 0f32;
8201                        for g in 0..ng {
8202                            let d = dot_i8_i8(
8203                                &buf[g * GROUP_SIZE..(g + 1) * GROUP_SIZE],
8204                                &act.xq[g * GROUP_SIZE..(g + 1) * GROUP_SIZE],
8205                            ) as f32
8206                                * act.sx;
8207                            dot += d * gscale(r, g);
8208                        }
8209                        for &(j, xv) in &act.outliers {
8210                            dot += (buf[j] as i8) as f32 * gscale(r, j / GROUP_SIZE) * xv;
8211                        }
8212                        // SAFETY: disjoint (bi, r) cells per worker range.
8213                        unsafe { *out_addr.at(bi * rows + r) = dot };
8214                        bi += 1;
8215                    }
8216                });
8217            }
8218        };
8219        dispatch_rows(pool, rows, &run);
8220        return;
8221    }
8222
8223    let out_addr = SendMut(out.as_mut_ptr());
8224    let run = move |start: usize, end: usize| {
8225        ROW_F32.with(|rb| {
8226            let mut buf = rb.borrow_mut();
8227            buf.resize(cols, 0.0);
8228            for r in start..end {
8229                decode_f32(r, &mut buf);
8230                for bi in 0..b {
8231                    let x = &xs_all[bi * cols..(bi + 1) * cols];
8232                    let mut dot = 0f32;
8233                    for g in 0..ng {
8234                        let mut gd = 0f32;
8235                        for k in 0..GROUP_SIZE {
8236                            gd += buf[g * GROUP_SIZE + k] * x[g * GROUP_SIZE + k];
8237                        }
8238                        dot += gd * gscale(r, g);
8239                    }
8240                    // SAFETY: disjoint (bi, r) cells per worker range.
8241                    unsafe { *out_addr.at(bi * rows + r) = dot };
8242                }
8243            }
8244        })
8245    };
8246    dispatch_rows(pool, rows, &run);
8247}
8248
8249/// Build a GPU batch job for a q8-family mapped tensor (primary
8250/// shard): prescaled input + directory coordinates. None → not
8251/// GPU-eligible, caller stays on the CPU.
8252pub(crate) fn gpu_batch_job<'a>(
8253    t: &'a QTensor,
8254    x: &[f32],
8255) -> Option<(std::sync::Arc<CmfModel>, crate::gpu::BatchJob<'a>)> {
8256    match t {
8257        QTensor::Mapped {
8258            model,
8259            idx,
8260            dtype: dt @ (TensorDtype::Q8Row | TensorDtype::Q8_2f),
8261            rows,
8262            cols,
8263            row_scale,
8264            col_field,
8265            ..
8266        } => Some((
8267            model.clone(),
8268            crate::gpu::BatchJob {
8269                idx: *idx,
8270                rows: *rows,
8271                cols: *cols,
8272                row_scale,
8273                xs: prescale(x, col_field, *dt).into_owned(),
8274                layout: crate::gpu::BatchLayout::Q8,
8275            },
8276        )),
8277        // q1: raw f32 activations, tile-embedded scales.
8278        QTensor::Mapped {
8279            model,
8280            idx,
8281            dtype: TensorDtype::Q1,
8282            rows,
8283            cols,
8284            ..
8285        } => Some((
8286            model.clone(),
8287            crate::gpu::BatchJob {
8288                idx: *idx,
8289                rows: *rows,
8290                cols: *cols,
8291                row_scale: &[],
8292                xs: x.to_vec(),
8293                layout: crate::gpu::BatchLayout::Q1,
8294            },
8295        )),
8296        // q4_tiled / q4tp: raw f32 activations; the scales live in the
8297        // payload (inline tiles / row ladder), so row_scale stays empty.
8298        // The GDN projection batch already runs these layouts on Metal —
8299        // this arm lets the attention QKV batch reach the same kernels.
8300        QTensor::Mapped {
8301            model,
8302            idx,
8303            dtype: dt @ (TensorDtype::Q4Tiled | TensorDtype::Q4TiledP),
8304            rows,
8305            cols,
8306            ..
8307        } => Some((
8308            model.clone(),
8309            crate::gpu::BatchJob {
8310                idx: *idx,
8311                rows: *rows,
8312                cols: *cols,
8313                row_scale: &[],
8314                xs: x.to_vec(),
8315                layout: if *dt == TensorDtype::Q4Tiled {
8316                    crate::gpu::BatchLayout::Q4t
8317                } else {
8318                    crate::gpu::BatchLayout::Q4tp
8319                },
8320            },
8321        )),
8322        _ => None,
8323    }
8324}
8325
8326thread_local! {
8327    static PRESCALE_BUF1: std::cell::RefCell<Vec<f32>> = const { std::cell::RefCell::new(Vec::new()) };
8328    static PRESCALE_BUF2: std::cell::RefCell<Vec<f32>> = const { std::cell::RefCell::new(Vec::new()) };
8329}
8330
8331pub(crate) fn prescale<'a>(
8332    x: &'a [f32],
8333    col_field: &[f32],
8334    dtype: TensorDtype,
8335) -> std::borrow::Cow<'a, [f32]> {
8336    if dtype == TensorDtype::Q8_2f {
8337        x.iter().zip(col_field).map(|(a, c)| a * c).collect()
8338    } else {
8339        std::borrow::Cow::Borrowed(x)
8340    }
8341}
8342
8343/// θ col-field fold for q8_2f activations. Borrowed pass-through for
8344/// every other dtype, using thread-local buffers to eliminate per-matvec allocations.
8345pub(crate) fn prescale_with<R, F: FnOnce(&[f32]) -> R>(
8346    x: &[f32],
8347    col_field: &[f32],
8348    dtype: TensorDtype,
8349    buf_id: u8,
8350    f: F,
8351) -> R {
8352    if dtype == TensorDtype::Q8_2f {
8353        if buf_id == 1 {
8354            PRESCALE_BUF1.with(|b| {
8355                let mut buf = b.borrow_mut();
8356                buf.clear();
8357                buf.extend(x.iter().zip(col_field).map(|(a, c)| a * c));
8358                f(&buf)
8359            })
8360        } else {
8361            PRESCALE_BUF2.with(|b| {
8362                let mut buf = b.borrow_mut();
8363                buf.clear();
8364                buf.extend(x.iter().zip(col_field).map(|(a, c)| a * c));
8365                f(&buf)
8366            })
8367        }
8368    } else {
8369        f(x)
8370    }
8371}
8372
8373// ───────────────────── x86-64 AVX2 kernels (roadmap этап 2) ─────────────────────
8374
8375/// AVX2+FMA available? Default ON when the CPU supports both;
8376/// `CMF_AVX2=0` disables (falls back to the autovectorized loops).
8377#[cfg(target_arch = "x86_64")]
8378pub(crate) fn avx2_enabled() -> bool {
8379    use std::sync::OnceLock;
8380    static ON: OnceLock<bool> = OnceLock::new();
8381    *ON.get_or_init(|| {
8382        std::env::var("CMF_AVX2").map(|v| v != "0").unwrap_or(true)
8383            && std::arch::is_x86_feature_detected!("avx2")
8384            && std::arch::is_x86_feature_detected!("fma")
8385    })
8386}
8387
8388/// AVX2 A8W8 allowed? The quantized-activation contract is switched by
8389/// the SAME env as the ARM SDOT path: `CMF_SDOT=0` keeps exact kernels
8390/// (the golden-parity exact gate relies on it) — AVX2 f32 kernels stay
8391/// active either way, they are exact (regrouped sums only).
8392#[cfg(target_arch = "x86_64")]
8393fn avx2_a8w8_enabled() -> bool {
8394    use std::sync::OnceLock;
8395    static ON: OnceLock<bool> = OnceLock::new();
8396    *ON.get_or_init(|| {
8397        avx2_enabled() && std::env::var("CMF_SDOT").map(|v| v != "0").unwrap_or(true)
8398    })
8399}
8400
8401/// A8W8 quantized-activation path available on THIS machine? One
8402/// switch across architectures: ARM dotprod (CMF_SDOT) or x86 AVX2
8403/// (CMF_AVX2 + the same CMF_SDOT exact-contract override).
8404#[inline]
8405pub(crate) fn a8w8_enabled() -> bool {
8406    #[cfg(target_arch = "aarch64")]
8407    {
8408        sdot_enabled()
8409    }
8410    #[cfg(target_arch = "x86_64")]
8411    {
8412        avx2_a8w8_enabled()
8413    }
8414    #[cfg(not(any(target_arch = "aarch64", target_arch = "x86_64")))]
8415    {
8416        false
8417    }
8418}
8419
8420/// int8·int8 dot dispatch: SDOT on ARM; AVX-512 VNNI (vpdpbusd) or AVX2
8421/// maddubs on x86. Callers are gated by `a8w8_enabled()`.
8422#[inline]
8423#[allow(unreachable_code)]
8424fn dot_i8_i8(w: &[u8], xq: &[i8]) -> i32 {
8425    #[cfg(target_arch = "aarch64")]
8426    unsafe {
8427        return dot_i8_sdot(w, xq);
8428    }
8429    #[cfg(target_arch = "x86_64")]
8430    unsafe {
8431        if avx512vnni_enabled() {
8432            return dot_i8_i8_vnni(w, xq);
8433        }
8434        return dot_i8_i8_avx2(w, xq);
8435    }
8436    w.iter()
8437        .zip(xq)
8438        .map(|(&a, &b)| (a as i8) as i32 * b as i32)
8439        .sum()
8440}
8441
8442/// AVX-512 VNNI available? (F+BW+VL+VNNI; `CMF_AVX512=0` falls back to
8443/// AVX2.) VL matters: short 32-byte groups (q4/vbit) ride the 256-bit
8444/// `vpdpbusd` encoding.
8445#[cfg(target_arch = "x86_64")]
8446fn avx512vnni_enabled() -> bool {
8447    use std::sync::OnceLock;
8448    static ON: OnceLock<bool> = OnceLock::new();
8449    *ON.get_or_init(|| {
8450        std::env::var("CMF_AVX512")
8451            .map(|v| v != "0")
8452            .unwrap_or(true)
8453            && std::arch::is_x86_feature_detected!("avx512f")
8454            && std::arch::is_x86_feature_detected!("avx512bw")
8455            && std::arch::is_x86_feature_detected!("avx512vl")
8456            && std::arch::is_x86_feature_detected!("avx512vnni")
8457    })
8458}
8459
8460/// Grouped-codec VNNI arms (the q4t/q4b/q1/q1t tile kernels): default
8461/// ON where AVX-512 VNNI exists (`CMF_VNNI_TILES=0` opt-out). Measured
8462/// on Ryzen 7950X (Zen4, 3 alternating process pairs, blocked GEMM
8463/// 4864×896 b=256): q4t 63→68 GF/s (+8%), q1 53→56 (+6%), q4b 72→75
8464/// (+4%) — consistent, no leg regressed. The tile kernels keep a
8465/// horizontal reduce per 32-weight group, so the `vpdpbusd` saving is
8466/// smaller than the long-dot q8 win (+13%), but it is real and free.
8467#[cfg(target_arch = "x86_64")]
8468fn vnni_tiles_enabled() -> bool {
8469    use std::sync::OnceLock;
8470    static ON: OnceLock<bool> = OnceLock::new();
8471    *ON.get_or_init(|| {
8472        std::env::var("CMF_VNNI_TILES")
8473            .map(|v| v != "0")
8474            .unwrap_or(true)
8475            && avx512vnni_enabled()
8476    })
8477}
8478
8479/// One 256-bit u8×i8 dot → i32 via `vpdpbusd` into a fresh accumulator
8480/// plus the same horizontal reduce the AVX2 kernels use. Products are
8481/// bounded (|w| ≤ 8 or ≤ 1), so maddubs never saturated — the i32 sum
8482/// is bit-identical to the maddubs+madd pair it replaces.
8483#[cfg(target_arch = "x86_64")]
8484#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
8485#[inline]
8486unsafe fn dpbusd_hsum(aw: core::arch::x86_64::__m256i, xs: core::arch::x86_64::__m256i) -> i32 {
8487    // SAFETY: pure register math.
8488    unsafe {
8489        use core::arch::x86_64::*;
8490        let d = _mm256_dpbusd_epi32(_mm256_setzero_si256(), aw, xs);
8491        let hi128 = _mm256_extracti128_si256::<1>(d);
8492        let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
8493        let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
8494        let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
8495        _mm_cvtsi128_si32(s32)
8496    }
8497}
8498
8499/// int8·int8 via AVX-512 VNNI: `vpdpbusd` fuses the maddubs+madd+add
8500/// triple into one u8×i8 dot-accumulate. AVX-512 has no vpsignb, so the
8501/// |w|·sign(x,w) trick becomes |w| × (x negated where w<0) via a mask
8502/// subtract — w==0 lanes contribute 0 through |w|=0 either way.
8503#[cfg(target_arch = "x86_64")]
8504#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
8505unsafe fn dot_i8_i8_vnni(w: &[u8], xq: &[i8]) -> i32 {
8506    // SAFETY: callers uphold slice-length contracts (see call sites).
8507    unsafe {
8508        use core::arch::x86_64::*;
8509        let n = w.len();
8510        let mut j = 0usize;
8511        let mut total: i32;
8512        // 4 independent accumulators: vpdpbusd is its own loop-carried
8513        // dependency (~5-cycle latency) — a single-acc loop runs
8514        // latency-bound and LOSES to the AVX2 maddubs kernel, measured
8515        // on Granite Rapids.
8516        {
8517            #[inline(always)]
8518            unsafe fn step(
8519                w: *const u8,
8520                x: *const i8,
8521                acc: core::arch::x86_64::__m512i,
8522            ) -> core::arch::x86_64::__m512i {
8523                unsafe {
8524                    use core::arch::x86_64::*;
8525                    let wv = _mm512_loadu_si512(w as *const _);
8526                    let xv = _mm512_loadu_si512(x as *const _);
8527                    let aw = _mm512_abs_epi8(wv);
8528                    let neg = _mm512_movepi8_mask(wv);
8529                    let sx = _mm512_mask_sub_epi8(xv, neg, _mm512_setzero_si512(), xv);
8530                    _mm512_dpbusd_epi32(acc, aw, sx)
8531                }
8532            }
8533            let (mut a0, mut a1, mut a2, mut a3) = (
8534                _mm512_setzero_si512(),
8535                _mm512_setzero_si512(),
8536                _mm512_setzero_si512(),
8537                _mm512_setzero_si512(),
8538            );
8539            while j + 256 <= n {
8540                a0 = step(w.as_ptr().add(j), xq.as_ptr().add(j), a0);
8541                a1 = step(w.as_ptr().add(j + 64), xq.as_ptr().add(j + 64), a1);
8542                a2 = step(w.as_ptr().add(j + 128), xq.as_ptr().add(j + 128), a2);
8543                a3 = step(w.as_ptr().add(j + 192), xq.as_ptr().add(j + 192), a3);
8544                j += 256;
8545            }
8546            while j + 64 <= n {
8547                a0 = step(w.as_ptr().add(j), xq.as_ptr().add(j), a0);
8548                j += 64;
8549            }
8550            let s01 = _mm512_add_epi32(a0, a1);
8551            let s23 = _mm512_add_epi32(a2, a3);
8552            total = _mm512_reduce_add_epi32(_mm512_add_epi32(s01, s23));
8553        }
8554        // 32-wide (q4/vbit groups are exactly 32 bytes).
8555        if j + 32 <= n {
8556            let wv = _mm256_loadu_si256(w.as_ptr().add(j) as *const __m256i);
8557            let xv = _mm256_loadu_si256(xq.as_ptr().add(j) as *const __m256i);
8558            let d = _mm256_dpbusd_epi32(
8559                _mm256_setzero_si256(),
8560                _mm256_abs_epi8(wv),
8561                _mm256_sign_epi8(xv, wv),
8562            );
8563            let hi128 = _mm256_extracti128_si256::<1>(d);
8564            let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
8565            let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
8566            let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
8567            total += _mm_cvtsi128_si32(s32);
8568            j += 32;
8569        }
8570        while j < n {
8571            total += (w[j] as i8) as i32 * xq[j] as i32;
8572            j += 1;
8573        }
8574        total
8575    }
8576}
8577
8578/// i8 row · f32 x via AVX2/FMA (x86 mirror of `dot_i8_f32_neon`).
8579#[cfg(target_arch = "x86_64")]
8580#[target_feature(enable = "avx2,fma")]
8581unsafe fn dot_i8_f32_avx2(w: &[u8], x: &[f32]) -> f32 {
8582    // SAFETY: callers uphold slice-length contracts (see call sites).
8583    unsafe {
8584        use core::arch::x86_64::*;
8585        let n = x.len();
8586        let wp = w.as_ptr();
8587        let xp = x.as_ptr();
8588        let (mut a0, mut a1) = (_mm256_setzero_ps(), _mm256_setzero_ps());
8589        let mut j = 0usize;
8590        while j + 16 <= n {
8591            let wb = _mm_loadu_si128(wp.add(j) as *const __m128i);
8592            let lo = _mm256_cvtepi8_epi32(wb);
8593            let hi = _mm256_cvtepi8_epi32(_mm_srli_si128::<8>(wb));
8594            a0 = _mm256_fmadd_ps(_mm256_cvtepi32_ps(lo), _mm256_loadu_ps(xp.add(j)), a0);
8595            a1 = _mm256_fmadd_ps(_mm256_cvtepi32_ps(hi), _mm256_loadu_ps(xp.add(j + 8)), a1);
8596            j += 16;
8597        }
8598        let acc = _mm256_add_ps(a0, a1);
8599        let hi128 = _mm256_extractf128_ps::<1>(acc);
8600        let s128 = _mm_add_ps(_mm256_castps256_ps128(acc), hi128);
8601        let s64 = _mm_add_ps(s128, _mm_movehl_ps(s128, s128));
8602        let s32 = _mm_add_ss(s64, _mm_shuffle_ps::<1>(s64, s64));
8603        let mut sum = _mm_cvtss_f32(s32);
8604        while j < n {
8605            sum += (*wp.add(j) as i8) as f32 * *xp.add(j);
8606            j += 1;
8607        }
8608        sum
8609    }
8610}
8611
8612/// int8(weight)·int8(activation) → i32 via AVX2 maddubs — the x86
8613/// analogue of the SDOT A8W8 path. `maddubs` takes u8×i8, so the
8614/// standard sign trick applies: |w| × sign(x, w) ≡ w × x per lane.
8615/// Pair saturation is safe: |w|≤128, |x|≤127 → 2·128·127 < 32767.
8616#[cfg(target_arch = "x86_64")]
8617#[target_feature(enable = "avx2")]
8618unsafe fn dot_i8_i8_avx2(w: &[u8], xq: &[i8]) -> i32 {
8619    // SAFETY: callers uphold slice-length contracts (see call sites).
8620    unsafe {
8621        use core::arch::x86_64::*;
8622        let n = w.len();
8623        let ones = _mm256_set1_epi16(1);
8624        let mut acc = _mm256_setzero_si256();
8625        let mut j = 0usize;
8626        while j + 32 <= n {
8627            let wv = _mm256_loadu_si256(w.as_ptr().add(j) as *const __m256i);
8628            let xv = _mm256_loadu_si256(xq.as_ptr().add(j) as *const __m256i);
8629            let p16 = _mm256_maddubs_epi16(_mm256_abs_epi8(wv), _mm256_sign_epi8(xv, wv));
8630            acc = _mm256_add_epi32(acc, _mm256_madd_epi16(p16, ones));
8631            j += 32;
8632        }
8633        let hi128 = _mm256_extracti128_si256::<1>(acc);
8634        let s128 = _mm_add_epi32(_mm256_castsi256_si128(acc), hi128);
8635        let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
8636        let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
8637        let mut s = _mm_cvtsi128_si32(s32);
8638        while j < n {
8639            s += (w[j] as i8) as i32 * xq[j] as i32;
8640            j += 1;
8641        }
8642        s
8643    }
8644}
8645
8646/// smmla 2×4: one instruction covers a 2-row × 2-activation × 8-deep
8647/// tile (32 MACs vs sdot's 16) — the weight pair loads once per 8-k
8648/// slice as a combined 2×8 register and meets two activation pairs.
8649#[cfg(target_arch = "aarch64")]
8650#[target_feature(enable = "neon,i8mm")]
8651unsafe fn dot_i8_smmla_2x4(w0: &[u8], w1: &[u8], xs: [&[i8]; 4]) -> [[i32; 4]; 2] {
8652    // SAFETY: callers uphold slice-length contracts.
8653    unsafe {
8654        use core::arch::aarch64::*;
8655        use core::arch::asm;
8656        let n = w0.len();
8657        let w0p = w0.as_ptr() as *const i8;
8658        let w1p = w1.as_ptr() as *const i8;
8659        // acc01 holds [c(r0,x0) c(r0,x1) c(r1,x0) c(r1,x1)]; acc23 the
8660        // same for x2/x3.
8661        let mut acc01 = vdupq_n_s32(0);
8662        let mut acc23 = vdupq_n_s32(0);
8663        let mut i = 0usize;
8664        while i + 8 <= n {
8665            let wa = vcombine_s8(vld1_s8(w0p.add(i)), vld1_s8(w1p.add(i)));
8666            let xb01 = vcombine_s8(
8667                vld1_s8(xs[0].as_ptr().add(i)),
8668                vld1_s8(xs[1].as_ptr().add(i)),
8669            );
8670            let xb23 = vcombine_s8(
8671                vld1_s8(xs[2].as_ptr().add(i)),
8672                vld1_s8(xs[3].as_ptr().add(i)),
8673            );
8674            asm!(
8675                "smmla {a01:v}.4s, {w:v}.16b, {x01:v}.16b",
8676                "smmla {a23:v}.4s, {w:v}.16b, {x23:v}.16b",
8677                a01 = inout(vreg) acc01, a23 = inout(vreg) acc23,
8678                w = in(vreg) wa, x01 = in(vreg) xb01, x23 = in(vreg) xb23,
8679                options(pure, nomem, nostack),
8680            );
8681            i += 8;
8682        }
8683        let mut out = [[0i32; 4]; 2];
8684        let a01: [i32; 4] = core::mem::transmute(acc01);
8685        let a23: [i32; 4] = core::mem::transmute(acc23);
8686        out[0][0] = a01[0];
8687        out[0][1] = a01[1];
8688        out[1][0] = a01[2];
8689        out[1][1] = a01[3];
8690        out[0][2] = a23[0];
8691        out[0][3] = a23[1];
8692        out[1][2] = a23[2];
8693        out[1][3] = a23[3];
8694        if i < n {
8695            for (k, x) in xs.iter().enumerate() {
8696                for j in i..n {
8697                    out[0][k] += (w0[j] as i8) as i32 * x[j] as i32;
8698                    out[1][k] += (w1[j] as i8) as i32 * x[j] as i32;
8699                }
8700            }
8701        }
8702        out
8703    }
8704}
8705
8706/// ARM twin of the x86 blocked prefill GEMM: two weight rows stay in
8707/// registers across four activation streams, eight sdot accumulators.
8708/// (The per-row form re-read each W row once per activation.)
8709#[cfg(target_arch = "aarch64")]
8710#[target_feature(enable = "neon,dotprod")]
8711unsafe fn dot_i8_sdot_2x4(w0: &[u8], w1: &[u8], xs: [&[i8]; 4]) -> [[i32; 4]; 2] {
8712    // SAFETY: callers uphold slice-length contracts.
8713    unsafe {
8714        use core::arch::aarch64::*;
8715        use core::arch::asm;
8716        let n = w0.len();
8717        let w0p = w0.as_ptr() as *const i8;
8718        let w1p = w1.as_ptr() as *const i8;
8719        let mut acc = [[vdupq_n_s32(0); 4]; 2];
8720        let mut i = 0usize;
8721        while i + 16 <= n {
8722            let wv0 = vld1q_s8(w0p.add(i));
8723            let wv1 = vld1q_s8(w1p.add(i));
8724            for (k, x) in xs.iter().enumerate() {
8725                let xv = vld1q_s8(x.as_ptr().add(i));
8726                let (mut a0, mut a1) = (acc[0][k], acc[1][k]);
8727                asm!(
8728                    "sdot {a0:v}.4s, {w0:v}.16b, {x:v}.16b",
8729                    "sdot {a1:v}.4s, {w1:v}.16b, {x:v}.16b",
8730                    a0 = inout(vreg) a0, a1 = inout(vreg) a1,
8731                    w0 = in(vreg) wv0, w1 = in(vreg) wv1, x = in(vreg) xv,
8732                    options(pure, nomem, nostack),
8733                );
8734                acc[0][k] = a0;
8735                acc[1][k] = a1;
8736            }
8737            i += 16;
8738        }
8739        let mut out = [[0i32; 4]; 2];
8740        for r in 0..2 {
8741            for k in 0..4 {
8742                out[r][k] = vaddvq_s32(acc[r][k]);
8743            }
8744        }
8745        if i < n {
8746            for (k, x) in xs.iter().enumerate() {
8747                for j in i..n {
8748                    out[0][k] += (w0[j] as i8) as i32 * x[j] as i32;
8749                    out[1][k] += (w1[j] as i8) as i32 * x[j] as i32;
8750                }
8751            }
8752        }
8753        out
8754    }
8755}
8756
8757/// Blocked 2 weight rows × 4 activations for the prefill GEMM
8758/// (roadmap P0: packed panels + multi-row accumulators). The two rows'
8759/// abs() live in registers across all four activation streams; the
8760/// sign-fixup is recomputed per pair (the price of the maddubs trick).
8761/// Returns raw i8·i8 dots; the caller applies scales and outliers.
8762#[cfg(target_arch = "x86_64")]
8763#[target_feature(enable = "avx2")]
8764unsafe fn dot_i8_i8_avx2_2x4(w0: &[u8], w1: &[u8], xs: [&[i8]; 4]) -> [[i32; 4]; 2] {
8765    // SAFETY: callers uphold slice-length contracts.
8766    unsafe {
8767        use core::arch::x86_64::*;
8768        let n = w0.len();
8769        let ones = _mm256_set1_epi16(1);
8770        let mut acc = [[_mm256_setzero_si256(); 4]; 2];
8771        let mut j = 0usize;
8772        while j + 32 <= n {
8773            let wv0 = _mm256_loadu_si256(w0.as_ptr().add(j) as *const __m256i);
8774            let wv1 = _mm256_loadu_si256(w1.as_ptr().add(j) as *const __m256i);
8775            let aw0 = _mm256_abs_epi8(wv0);
8776            let aw1 = _mm256_abs_epi8(wv1);
8777            for (k, x) in xs.iter().enumerate() {
8778                let xv = _mm256_loadu_si256(x.as_ptr().add(j) as *const __m256i);
8779                let p0 = _mm256_maddubs_epi16(aw0, _mm256_sign_epi8(xv, wv0));
8780                acc[0][k] = _mm256_add_epi32(acc[0][k], _mm256_madd_epi16(p0, ones));
8781                let p1 = _mm256_maddubs_epi16(aw1, _mm256_sign_epi8(xv, wv1));
8782                acc[1][k] = _mm256_add_epi32(acc[1][k], _mm256_madd_epi16(p1, ones));
8783            }
8784            j += 32;
8785        }
8786        let mut out = [[0i32; 4]; 2];
8787        for r in 0..2 {
8788            for k in 0..4 {
8789                let a = acc[r][k];
8790                let hi128 = _mm256_extracti128_si256::<1>(a);
8791                let s128 = _mm_add_epi32(_mm256_castsi256_si128(a), hi128);
8792                let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
8793                let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
8794                out[r][k] = _mm_cvtsi128_si32(s32);
8795            }
8796        }
8797        if j < n {
8798            for (k, x) in xs.iter().enumerate() {
8799                for i in j..n {
8800                    out[0][k] += (w0[i] as i8) as i32 * x[i] as i32;
8801                    out[1][k] += (w1[i] as i8) as i32 * x[i] as i32;
8802                }
8803            }
8804        }
8805        out
8806    }
8807}
8808
8809/// AVX2/VNNI q8 row dot with exact outlier correction (x86 mirror of
8810/// `row_dot_sdot` — same A8W8 contract). With AVX-512 VNNI the row goes
8811/// through the bias trick: Σ(w+128)·x via pure `vpdpbusd` (no per-lane
8812/// sign fixups), corrected by −128·Σx with Σx precomputed per split.
8813#[cfg(target_arch = "x86_64")]
8814#[inline]
8815fn row_dot_avx2(row: &[u8], act: &SplitAct) -> f32 {
8816    let dot = if avx512vnni_enabled() && row.len() >= 64 {
8817        (unsafe { dot_u8p128_i8_vnni(row, &act.xq) }) - 128 * act.xsum
8818    } else {
8819        unsafe { dot_i8_i8_avx2(row, &act.xq) }
8820    };
8821    let mut acc = dot as f32 * act.sx;
8822    for &(j, xv) in &act.outliers {
8823        acc += (row[j] as i8) as f32 * xv;
8824    }
8825    acc
8826}
8827
8828/// Σ (w[i]+128)·x[i] via pure `vpdpbusd` — the caller subtracts
8829/// 128·Σx. Four independent accumulators (dpbusd is ~5-cycle latency;
8830/// a single-acc loop runs latency-bound, measured on Granite Rapids).
8831#[cfg(target_arch = "x86_64")]
8832#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
8833unsafe fn dot_u8p128_i8_vnni(w: &[u8], xq: &[i8]) -> i32 {
8834    // SAFETY: callers uphold slice-length contracts (see call sites).
8835    unsafe {
8836        use core::arch::x86_64::*;
8837        let n = w.len();
8838        let flip = _mm512_set1_epi8(-128); // XOR 0x80: i8 w → u8 (w+128)
8839        #[inline(always)]
8840        unsafe fn step(
8841            w: *const u8,
8842            x: *const i8,
8843            flip: core::arch::x86_64::__m512i,
8844            acc: core::arch::x86_64::__m512i,
8845        ) -> core::arch::x86_64::__m512i {
8846            unsafe {
8847                use core::arch::x86_64::*;
8848                let wv = _mm512_xor_si512(_mm512_loadu_si512(w as *const _), flip);
8849                _mm512_dpbusd_epi32(acc, wv, _mm512_loadu_si512(x as *const _))
8850            }
8851        }
8852        let (mut a0, mut a1, mut a2, mut a3) = (
8853            _mm512_setzero_si512(),
8854            _mm512_setzero_si512(),
8855            _mm512_setzero_si512(),
8856            _mm512_setzero_si512(),
8857        );
8858        let mut j = 0usize;
8859        while j + 256 <= n {
8860            a0 = step(w.as_ptr().add(j), xq.as_ptr().add(j), flip, a0);
8861            a1 = step(w.as_ptr().add(j + 64), xq.as_ptr().add(j + 64), flip, a1);
8862            a2 = step(w.as_ptr().add(j + 128), xq.as_ptr().add(j + 128), flip, a2);
8863            a3 = step(w.as_ptr().add(j + 192), xq.as_ptr().add(j + 192), flip, a3);
8864            j += 256;
8865        }
8866        while j + 64 <= n {
8867            a0 = step(w.as_ptr().add(j), xq.as_ptr().add(j), flip, a0);
8868            j += 64;
8869        }
8870        let mut total = _mm512_reduce_add_epi32(_mm512_add_epi32(
8871            _mm512_add_epi32(a0, a1),
8872            _mm512_add_epi32(a2, a3),
8873        ));
8874        // Scalar tail: (w as i8) + 128 ≡ (w as u8) ^ 0x80.
8875        while j < n {
8876            total += ((w[j] ^ 0x80) as i32) * xq[j] as i32;
8877            j += 1;
8878        }
8879        total
8880    }
8881}
8882
8883/// One q4 row via AVX2: nibbles → centered i8 (unpacklo/hi restores the
8884/// writer's flat order, same as the NEON vzip pair), maddubs against
8885/// the pre-quantized activation group, × the group's f16 scale. Pair
8886/// saturation safe: |w|≤8, |x|≤127 → 2·8·127 ≪ 32767. Mirror of
8887/// `dot_q4_row_sdot`.
8888#[cfg(target_arch = "x86_64")]
8889#[target_feature(enable = "avx2")]
8890unsafe fn dot_q4_row_avx2(packed: &[u8], scales: &[u8], g0: usize, gpr: usize, xq: &[i8]) -> f32 {
8891    // SAFETY: callers uphold slice-length contracts (16 packed bytes and
8892    // 2 scale bytes per group; xq.len() == gpr·GROUP_SIZE).
8893    unsafe {
8894        use core::arch::x86_64::*;
8895        let lomask = _mm_set1_epi8(0x0F);
8896        let eight = _mm256_set1_epi8(8);
8897        let ones = _mm256_set1_epi16(1);
8898        let mut acc = 0f32;
8899        for gi in 0..gpr {
8900            let g = g0 + gi;
8901            let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
8902            let b = _mm_loadu_si128(packed.as_ptr().add(g * 16) as *const __m128i);
8903            let lo = _mm_and_si128(b, lomask);
8904            let hi = _mm_and_si128(_mm_srli_epi16::<4>(b), lomask);
8905            let w = _mm256_sub_epi8(
8906                _mm256_set_m128i(_mm_unpackhi_epi8(lo, hi), _mm_unpacklo_epi8(lo, hi)),
8907                eight,
8908            );
8909            let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
8910            let p16 = _mm256_maddubs_epi16(_mm256_abs_epi8(w), _mm256_sign_epi8(x, w));
8911            let d = _mm256_madd_epi16(p16, ones);
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            acc += _mm_cvtsi128_si32(s32) as f32 * s;
8917        }
8918        acc
8919    }
8920}
8921
8922/// Two-activation q4 row via AVX2: nibbles unpacked ONCE per group,
8923/// both activations dotted against the same centered i8 register.
8924#[cfg(target_arch = "x86_64")]
8925#[target_feature(enable = "avx2")]
8926unsafe fn dot_q4_row_avx2_2(
8927    packed: &[u8],
8928    scales: &[u8],
8929    g0: usize,
8930    gpr: usize,
8931    xq1: &[i8],
8932    xq2: &[i8],
8933) -> (f32, f32) {
8934    // SAFETY: callers uphold slice-length contracts (see dot_q4_row_avx2).
8935    unsafe {
8936        use core::arch::x86_64::*;
8937        let lomask = _mm_set1_epi8(0x0F);
8938        let eight = _mm256_set1_epi8(8);
8939        let ones = _mm256_set1_epi16(1);
8940        let (mut acc1, mut acc2) = (0f32, 0f32);
8941        #[inline(always)]
8942        unsafe fn hsum(d: core::arch::x86_64::__m256i) -> i32 {
8943            unsafe {
8944                use core::arch::x86_64::*;
8945                let hi128 = _mm256_extracti128_si256::<1>(d);
8946                let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
8947                let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
8948                let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
8949                _mm_cvtsi128_si32(s32)
8950            }
8951        }
8952        for gi in 0..gpr {
8953            let g = g0 + gi;
8954            let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
8955            let b = _mm_loadu_si128(packed.as_ptr().add(g * 16) as *const __m128i);
8956            let lo = _mm_and_si128(b, lomask);
8957            let hi = _mm_and_si128(_mm_srli_epi16::<4>(b), lomask);
8958            let w = _mm256_sub_epi8(
8959                _mm256_set_m128i(_mm_unpackhi_epi8(lo, hi), _mm_unpacklo_epi8(lo, hi)),
8960                eight,
8961            );
8962            let aw = _mm256_abs_epi8(w);
8963            let x1 = _mm256_loadu_si256(xq1.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
8964            let x2 = _mm256_loadu_si256(xq2.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
8965            let d1 = _mm256_madd_epi16(_mm256_maddubs_epi16(aw, _mm256_sign_epi8(x1, w)), ones);
8966            let d2 = _mm256_madd_epi16(_mm256_maddubs_epi16(aw, _mm256_sign_epi8(x2, w)), ones);
8967            acc1 += hsum(d1) as f32 * s;
8968            acc2 += hsum(d2) as f32 * s;
8969        }
8970        (acc1, acc2)
8971    }
8972}
8973
8974/// One q8 row range via AVX2 (x86 mirror of `q8_range_sdot`).
8975#[cfg(target_arch = "x86_64")]
8976fn q8_range_avx2(
8977    q: &[u8],
8978    row_scale: &[f32],
8979    act: &SplitAct,
8980    cols: usize,
8981    out_addr: SendMut,
8982    start: usize,
8983    end: usize,
8984) {
8985    for o in start..end {
8986        let v = row_dot_avx2(&q[o * cols..(o + 1) * cols], act) * row_scale[o];
8987        // SAFETY: disjoint row ranges per worker.
8988        unsafe { *out_addr.at(o) = v };
8989    }
8990}
8991
8992/// Two-input q8 row range via AVX2 (x86 mirror of `q8_range2_sdot`).
8993#[cfg(target_arch = "x86_64")]
8994#[allow(clippy::too_many_arguments)]
8995fn q8_range2_avx2(
8996    q: &[u8],
8997    row_scale: &[f32],
8998    a1: &SplitAct,
8999    a2: &SplitAct,
9000    cols: usize,
9001    p1: SendMut,
9002    p2: SendMut,
9003    start: usize,
9004    end: usize,
9005) {
9006    for o in start..end {
9007        let row = &q[o * cols..(o + 1) * cols];
9008        // SAFETY: disjoint row ranges per worker.
9009        unsafe {
9010            *p1.at(o) = row_dot_avx2(row, a1) * row_scale[o];
9011            *p2.at(o) = row_dot_avx2(row, a2) * row_scale[o];
9012        }
9013    }
9014}
9015
9016// ───────────────────── A8W8 SDOT path (port of vmfcore, ×1.78 decode) ─────────────────────
9017
9018/// ARMv8.6 i8mm (smmla): 32 int8 MACs per instruction vs sdot's 16 —
9019/// yet MEASURED 2.4× SLOWER than the blocked sdot on Apple silicon
9020/// (108 vs 264 GF/s): the on-the-fly vcombine packing and the two-
9021/// accumulator dependency chain swamp the MAC advantage, and Apple's
9022/// four SIMD pipes already keep sdot fed. OPT-IN (CMF_I8MM=1) for
9023/// field trials on Cortex-A710/X-class parts with two pipes, where the
9024/// balance may differ; a pre-interleaved weight layout (repack infra)
9025/// is the known path if it ever earns its keep.
9026#[cfg(target_arch = "aarch64")]
9027fn i8mm_enabled() -> bool {
9028    static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
9029    *ON.get_or_init(|| {
9030        std::env::var("CMF_I8MM").map(|v| v == "1").unwrap_or(false)
9031            && std::arch::is_aarch64_feature_detected!("i8mm")
9032    })
9033}
9034
9035/// SDOT enabled? Default ON when the CPU has ARMv8.2 dotprod;
9036/// `CMF_SDOT=0` disables (falls back to i8×f32 NEON).
9037/// (On non-ARM release builds only the test tolerance switch calls it.)
9038#[cfg_attr(not(target_arch = "aarch64"), allow(dead_code))]
9039fn sdot_enabled() -> bool {
9040    use std::sync::OnceLock;
9041    static ON: OnceLock<bool> = OnceLock::new();
9042    *ON.get_or_init(|| {
9043        let want = std::env::var("CMF_SDOT").map(|v| v != "0").unwrap_or(true);
9044        if !want {
9045            return false;
9046        }
9047
9048        #[cfg(target_arch = "aarch64")]
9049        {
9050            if std::arch::is_aarch64_feature_detected!("dotprod") {
9051                return true;
9052            }
9053            #[cfg(target_os = "android")]
9054            {
9055                if let Ok(cpuinfo) = std::fs::read_to_string("/proc/cpuinfo") {
9056                    if cpuinfo.lines().any(|l| {
9057                        (l.starts_with("Features") || l.starts_with("features"))
9058                            && l.contains("asimddp")
9059                    }) {
9060                        return true;
9061                    }
9062                }
9063            }
9064            false
9065        }
9066        #[cfg(not(target_arch = "aarch64"))]
9067        {
9068            false
9069        }
9070    })
9071}
9072
9073/// Two-field activation split (≡ vmfcore `q8_split_prep`): outlier
9074/// channels (>8·rms) are computed exactly in f32; the bulk (outliers
9075/// zeroed → clean absmax) goes through int8 SDOT. Computed ONCE per
9076/// matvec, shared by all rows/workers.
9077struct SplitAct {
9078    xq: Vec<i8>,
9079    sx: f32,
9080    outliers: Vec<(usize, f32)>,
9081    /// Σ xq — the VNNI bias-trick correction (`(w+128)·x` sums need
9082    /// `−128·Σx`); one i32 per split, computed once per matvec.
9083    #[cfg_attr(not(target_arch = "x86_64"), allow(dead_code))]
9084    xsum: i32,
9085}
9086
9087thread_local! {
9088    /// Recycled xq buffers: split_act runs for every matvec (~200/token)
9089    /// and its hidden-size allocation was steady-state heap churn.
9090    static XQ_FREE: std::cell::RefCell<Vec<Vec<i8>>> =
9091        const { std::cell::RefCell::new(Vec::new()) };
9092}
9093
9094impl Drop for SplitAct {
9095    fn drop(&mut self) {
9096        let buf = std::mem::take(&mut self.xq);
9097        if buf.capacity() > 0 {
9098            XQ_FREE.with(|f| {
9099                let mut f = f.borrow_mut();
9100                if f.len() < 16 {
9101                    f.push(buf);
9102                }
9103            });
9104        }
9105    }
9106}
9107
9108thread_local! {
9109    /// One scratch row per WORKER, kept for the life of the thread.
9110    ///
9111    /// The kernels take a row of group scales per dispatch, and a fresh
9112    /// `vec![0f32; gpr]` inside the closure is one allocation per worker per
9113    /// dispatch — on the release checkpoint about six thousand a token, a
9114    /// quarter of everything the benchmark counts.
9115    static KROW: std::cell::RefCell<Vec<f32>> = const { std::cell::RefCell::new(Vec::new()) };
9116}
9117
9118/// Borrow `n` floats of the calling worker's scratch. Nothing inside a
9119/// kernel body borrows it again, which is what keeps the RefCell honest.
9120#[inline]
9121fn with_krow<R>(n: usize, f: impl FnOnce(&mut [f32]) -> R) -> R {
9122    KROW.with(|s| {
9123        let mut b = s.borrow_mut();
9124        if b.len() < n {
9125            b.resize(n, 0.0);
9126        }
9127        f(&mut b[..n])
9128    })
9129}
9130
9131fn split_act(x: &[f32]) -> SplitAct {
9132    let n = x.len();
9133    let rms = (x.iter().map(|&v| (v * v) as f64).sum::<f64>() / n.max(1) as f64).sqrt() as f32;
9134    let thr = 8.0 * rms;
9135    // One pass: collect outliers and the bulk absmax (outliers excluded —
9136    // identical to the old zero-then-fold over a copied buffer, minus the
9137    // full-vector copy).
9138    let mut outliers: Vec<(usize, f32)> = Vec::new();
9139    let mut amax = 0f32;
9140    for (j, &v) in x.iter().enumerate() {
9141        let a = v.abs();
9142        if a > thr {
9143            outliers.push((j, v));
9144        } else if a > amax {
9145            amax = a;
9146        }
9147    }
9148    let sx = if amax > 0.0 { amax / 127.0 } else { 1.0 };
9149    let inv = 1.0 / sx;
9150    let mut xq = XQ_FREE.with(|f| f.borrow_mut().pop()).unwrap_or_default();
9151    xq.clear();
9152    xq.reserve(n);
9153    if outliers.is_empty() {
9154        xq.extend(
9155            x.iter()
9156                .map(|&v| (v * inv).round().clamp(-127.0, 127.0) as i8),
9157        );
9158    } else {
9159        // Outlier slots quantize to 0 (their exact term is added later).
9160        xq.extend(x.iter().map(|&v| {
9161            if v.abs() > thr {
9162                0
9163            } else {
9164                (v * inv).round().clamp(-127.0, 127.0) as i8
9165            }
9166        }));
9167    }
9168    let xsum = xq.iter().map(|&v| v as i32).sum();
9169    SplitAct {
9170        xq,
9171        sx,
9172        outliers,
9173        xsum,
9174    }
9175}
9176
9177fn split_act_q8_2f(x: &[f32], col: &[f32]) -> SplitAct {
9178    let n = x.len();
9179    let rms = (x
9180        .iter()
9181        .zip(col)
9182        .map(|(&a, &c)| {
9183            let v = a * c;
9184            (v * v) as f64
9185        })
9186        .sum::<f64>()
9187        / n.max(1) as f64)
9188        .sqrt() as f32;
9189    let thr = 8.0 * rms;
9190
9191    let mut outliers = Vec::new();
9192    let mut amax = 0f32;
9193    for (j, (&a, &c)) in x.iter().zip(col).enumerate() {
9194        let v = a * c;
9195        let s = v.abs();
9196        if s > thr {
9197            outliers.push((j, v));
9198        } else if s > amax {
9199            amax = s;
9200        }
9201    }
9202
9203    let sx = if amax > 0.0 { amax / 127.0 } else { 1.0 };
9204    let inv = 1.0 / sx;
9205    let mut xq = XQ_FREE.with(|f| f.borrow_mut().pop()).unwrap_or_default();
9206    xq.clear();
9207    xq.reserve(n);
9208    if outliers.is_empty() {
9209        xq.extend(
9210            x.iter()
9211                .zip(col)
9212                .map(|(&a, &c)| ((a * c) * inv).round().clamp(-127.0, 127.0) as i8),
9213        );
9214    } else {
9215        xq.extend(x.iter().zip(col).map(|(&a, &c)| {
9216            let v = a * c;
9217            if v.abs() > thr {
9218                0
9219            } else {
9220                (v * inv).round().clamp(-127.0, 127.0) as i8
9221            }
9222        }));
9223    }
9224    let xsum = xq.iter().map(|&v| v as i32).sum();
9225    SplitAct {
9226        xq,
9227        sx,
9228        outliers,
9229        xsum,
9230    }
9231}
9232
9233/// int8(weight)·int8(activation) → i32 via `sdot` (inline asm — the
9234/// vdotq intrinsic is unstable; port of vmfcore `dot_i8_sdot`).
9235#[cfg(target_arch = "aarch64")]
9236#[target_feature(enable = "neon,dotprod")]
9237unsafe fn dot_i8_sdot(w: &[u8], xq: &[i8]) -> i32 {
9238    // SAFETY: callers uphold slice-length contracts (see call sites).
9239    unsafe {
9240        use core::arch::aarch64::*;
9241        use core::arch::asm;
9242        let wp = w.as_ptr() as *const i8;
9243        let n = w.len();
9244        let (mut a0, mut a1, mut a2, mut a3) = (
9245            vdupq_n_s32(0),
9246            vdupq_n_s32(0),
9247            vdupq_n_s32(0),
9248            vdupq_n_s32(0),
9249        );
9250        let mut i = 0;
9251        while i + 64 <= n {
9252            let (w0, x0) = (vld1q_s8(wp.add(i)), vld1q_s8(xq.as_ptr().add(i)));
9253            let (w1, x1) = (vld1q_s8(wp.add(i + 16)), vld1q_s8(xq.as_ptr().add(i + 16)));
9254            let (w2, x2) = (vld1q_s8(wp.add(i + 32)), vld1q_s8(xq.as_ptr().add(i + 32)));
9255            let (w3, x3) = (vld1q_s8(wp.add(i + 48)), vld1q_s8(xq.as_ptr().add(i + 48)));
9256            asm!(
9257                "sdot {a0:v}.4s, {w0:v}.16b, {x0:v}.16b",
9258                "sdot {a1:v}.4s, {w1:v}.16b, {x1:v}.16b",
9259                "sdot {a2:v}.4s, {w2:v}.16b, {x2:v}.16b",
9260                "sdot {a3:v}.4s, {w3:v}.16b, {x3:v}.16b",
9261                a0 = inout(vreg) a0, a1 = inout(vreg) a1, a2 = inout(vreg) a2, a3 = inout(vreg) a3,
9262                w0 = in(vreg) w0, x0 = in(vreg) x0, w1 = in(vreg) w1, x1 = in(vreg) x1,
9263                w2 = in(vreg) w2, x2 = in(vreg) x2, w3 = in(vreg) w3, x3 = in(vreg) x3,
9264                options(pure, nomem, nostack),
9265            );
9266            i += 64;
9267        }
9268        while i + 16 <= n {
9269            let (wv, xv) = (vld1q_s8(wp.add(i)), vld1q_s8(xq.as_ptr().add(i)));
9270            asm!("sdot {a:v}.4s, {w:v}.16b, {x:v}.16b",
9271                 a = inout(vreg) a0, w = in(vreg) wv, x = in(vreg) xv, options(pure, nomem, nostack));
9272            i += 16;
9273        }
9274        let mut s = vaddvq_s32(vaddq_s32(vaddq_s32(a0, a1), vaddq_s32(a2, a3)));
9275        while i < n {
9276            s += (*wp.add(i)) as i32 * xq[i] as i32;
9277            i += 1;
9278        }
9279        s
9280    }
9281}
9282
9283/// Row-blocked SDOT: 4 output rows per pass — the activation chunk is
9284/// loaded once and reused, 4 independent accumulators hide sdot latency
9285/// (port of vmfcore `dot_i8_sdot_4rows`).
9286#[cfg(target_arch = "aarch64")]
9287#[target_feature(enable = "neon,dotprod")]
9288unsafe fn dot_i8_sdot_4rows(w0: &[u8], w1: &[u8], w2: &[u8], w3: &[u8], xq: &[i8]) -> [i32; 4] {
9289    // SAFETY: callers uphold slice-length contracts (see call sites).
9290    unsafe {
9291        use core::arch::aarch64::*;
9292        use core::arch::asm;
9293        let n = xq.len();
9294        let px = xq.as_ptr();
9295        let (p0, p1, p2, p3) = (
9296            w0.as_ptr() as *const i8,
9297            w1.as_ptr() as *const i8,
9298            w2.as_ptr() as *const i8,
9299            w3.as_ptr() as *const i8,
9300        );
9301        let (mut a0, mut a1, mut a2, mut a3) = (
9302            vdupq_n_s32(0),
9303            vdupq_n_s32(0),
9304            vdupq_n_s32(0),
9305            vdupq_n_s32(0),
9306        );
9307        let mut i = 0;
9308        while i + 16 <= n {
9309            let x = vld1q_s8(px.add(i));
9310            let v0 = vld1q_s8(p0.add(i));
9311            let v1 = vld1q_s8(p1.add(i));
9312            let v2 = vld1q_s8(p2.add(i));
9313            let v3 = vld1q_s8(p3.add(i));
9314            asm!(
9315                "sdot {a0:v}.4s, {v0:v}.16b, {x:v}.16b",
9316                "sdot {a1:v}.4s, {v1:v}.16b, {x:v}.16b",
9317                "sdot {a2:v}.4s, {v2:v}.16b, {x:v}.16b",
9318                "sdot {a3:v}.4s, {v3:v}.16b, {x:v}.16b",
9319                a0 = inout(vreg) a0, a1 = inout(vreg) a1, a2 = inout(vreg) a2, a3 = inout(vreg) a3,
9320                v0 = in(vreg) v0, v1 = in(vreg) v1, v2 = in(vreg) v2, v3 = in(vreg) v3, x = in(vreg) x,
9321                options(pure, nomem, nostack),
9322            );
9323            i += 16;
9324        }
9325        let mut r = [
9326            vaddvq_s32(a0),
9327            vaddvq_s32(a1),
9328            vaddvq_s32(a2),
9329            vaddvq_s32(a3),
9330        ];
9331        while i < n {
9332            let xi = *px.add(i) as i32;
9333            r[0] += (*p0.add(i)) as i32 * xi;
9334            r[1] += (*p1.add(i)) as i32 * xi;
9335            r[2] += (*p2.add(i)) as i32 * xi;
9336            r[3] += (*p3.add(i)) as i32 * xi;
9337            i += 1;
9338        }
9339        r
9340    }
9341}
9342
9343/// 4 interleaved rows in one pass: the repacked group is [r0[c], r1[c],
9344/// r2[c], r3[c]] per 16-byte chunk, so each iteration reads ONE 64-byte
9345/// line plus the shared activation chunk — a single sequential weight
9346/// stream per worker. Per-row accumulation is the same one-accumulator
9347/// scheme as `dot_i8_sdot_4rows`; integer sums are exact, so outputs
9348/// are bit-identical to the mmap-layout kernel.
9349#[cfg(target_arch = "aarch64")]
9350#[target_feature(enable = "neon,dotprod")]
9351unsafe fn dot_i8_sdot_4rows_il(g: &[u8], xq: &[i8]) -> [i32; 4] {
9352    // SAFETY: callers uphold slice-length contracts (g.len() == 4·n,
9353    // n % 16 == 0 — guaranteed by the repack gate).
9354    unsafe {
9355        use core::arch::aarch64::*;
9356        use core::arch::asm;
9357        let n = xq.len();
9358        let px = xq.as_ptr();
9359        let pg = g.as_ptr() as *const i8;
9360        let (mut a0, mut a1, mut a2, mut a3) = (
9361            vdupq_n_s32(0),
9362            vdupq_n_s32(0),
9363            vdupq_n_s32(0),
9364            vdupq_n_s32(0),
9365        );
9366        let mut i = 0;
9367        while i + 16 <= n {
9368            let x = vld1q_s8(px.add(i));
9369            let base = pg.add(4 * i);
9370            let v0 = vld1q_s8(base);
9371            let v1 = vld1q_s8(base.add(16));
9372            let v2 = vld1q_s8(base.add(32));
9373            let v3 = vld1q_s8(base.add(48));
9374            asm!(
9375                "sdot {a0:v}.4s, {v0:v}.16b, {x:v}.16b",
9376                "sdot {a1:v}.4s, {v1:v}.16b, {x:v}.16b",
9377                "sdot {a2:v}.4s, {v2:v}.16b, {x:v}.16b",
9378                "sdot {a3:v}.4s, {v3:v}.16b, {x:v}.16b",
9379                a0 = inout(vreg) a0, a1 = inout(vreg) a1, a2 = inout(vreg) a2, a3 = inout(vreg) a3,
9380                v0 = in(vreg) v0, v1 = in(vreg) v1, v2 = in(vreg) v2, v3 = in(vreg) v3, x = in(vreg) x,
9381                options(pure, nomem, nostack),
9382            );
9383            i += 16;
9384        }
9385        [
9386            vaddvq_s32(a0),
9387            vaddvq_s32(a1),
9388            vaddvq_s32(a2),
9389            vaddvq_s32(a3),
9390        ]
9391    }
9392}
9393
9394/// One q8 row range via SDOT (4-row blocks + tail) — the body of
9395/// `qmatvec`'s hot loop, extracted so multi-matrix jobs can drive the
9396/// SAME kernel for several tensors under one pool dispatch. `rep` — the
9397/// load-time interleaved repack (empty = mmap layout only); rows outside
9398/// full 4-row groups always come from the mmap layout.
9399#[cfg(target_arch = "aarch64")]
9400fn q8_range_sdot(
9401    q: &[u8],
9402    rep: &[u8],
9403    row_scale: &[f32],
9404    act: &SplitAct,
9405    cols: usize,
9406    out_addr: SendMut,
9407    start: usize,
9408    end: usize,
9409) {
9410    let mut o = start;
9411    // Leading rows to the group boundary (repack path only): the pool
9412    // splits row ranges arbitrarily, groups are absolute.
9413    if !rep.is_empty() {
9414        while o < end && o % 4 != 0 {
9415            let v = row_dot_sdot(&q[o * cols..(o + 1) * cols], act) * row_scale[o];
9416            unsafe { *out_addr.at(o) = v };
9417            o += 1;
9418        }
9419    }
9420    while o + 4 <= end {
9421        let r = if rep.is_empty() {
9422            unsafe {
9423                dot_i8_sdot_4rows(
9424                    &q[o * cols..(o + 1) * cols],
9425                    &q[(o + 1) * cols..(o + 2) * cols],
9426                    &q[(o + 2) * cols..(o + 3) * cols],
9427                    &q[(o + 3) * cols..(o + 4) * cols],
9428                    &act.xq,
9429                )
9430            }
9431        } else {
9432            unsafe { dot_i8_sdot_4rows_il(&rep[o * cols..(o + 4) * cols], &act.xq) }
9433        };
9434        for k in 0..4 {
9435            let mut acc = r[k] as f32 * act.sx;
9436            for &(j, xv) in &act.outliers {
9437                acc += (q[(o + k) * cols + j] as i8) as f32 * xv;
9438            }
9439            // SAFETY: disjoint row ranges per worker.
9440            unsafe { *out_addr.at(o + k) = acc * row_scale[o + k] };
9441        }
9442        o += 4;
9443    }
9444    while o < end {
9445        let v = row_dot_sdot(&q[o * cols..(o + 1) * cols], act) * row_scale[o];
9446        unsafe { *out_addr.at(o) = v };
9447        o += 1;
9448    }
9449}
9450
9451/// Two-input q8 row range via SDOT — `qmatvec2`'s hot loop, extracted
9452/// for the fused pair multi-matrix job (`matvec2_many`).
9453#[cfg(target_arch = "aarch64")]
9454#[allow(clippy::too_many_arguments)]
9455fn q8_range2_sdot(
9456    q: &[u8],
9457    row_scale: &[f32],
9458    a1: &SplitAct,
9459    a2: &SplitAct,
9460    cols: usize,
9461    p1: SendMut,
9462    p2: SendMut,
9463    start: usize,
9464    end: usize,
9465) {
9466    for o in start..end {
9467        let row = &q[o * cols..(o + 1) * cols];
9468        // SAFETY: disjoint row ranges per worker.
9469        unsafe {
9470            *p1.at(o) = row_dot_sdot(row, a1) * row_scale[o];
9471            *p2.at(o) = row_dot_sdot(row, a2) * row_scale[o];
9472        }
9473    }
9474}
9475
9476/// Two-input q8 row range, f32 kernel (non-SDOT) — same extraction.
9477#[allow(clippy::too_many_arguments)]
9478fn q8_range2_f32(
9479    q: &[u8],
9480    row_scale: &[f32],
9481    x1: &[f32],
9482    x2: &[f32],
9483    cols: usize,
9484    p1: SendMut,
9485    p2: SendMut,
9486    start: usize,
9487    end: usize,
9488) {
9489    for o in start..end {
9490        let row = &q[o * cols..(o + 1) * cols];
9491        // SAFETY: disjoint row ranges per worker.
9492        unsafe {
9493            *p1.at(o) = dot_i8_f32(row, x1) * row_scale[o];
9494            *p2.at(o) = dot_i8_f32(row, x2) * row_scale[o];
9495        }
9496    }
9497}
9498
9499/// Scalar/NEON-f32 q8 row range (non-SDOT platforms) — same extraction.
9500fn q8_range_f32(
9501    q: &[u8],
9502    row_scale: &[f32],
9503    xs: &[f32],
9504    cols: usize,
9505    out_addr: SendMut,
9506    start: usize,
9507    end: usize,
9508) {
9509    for o in start..end {
9510        let v = dot_i8_f32(&q[o * cols..(o + 1) * cols], xs) * row_scale[o];
9511        // SAFETY: disjoint row ranges per worker.
9512        unsafe { *out_addr.at(o) = v };
9513    }
9514}
9515
9516/// One q8 row against a split activation, portable: the per-arch fast
9517/// dots where they exist, the exact scalar loop elsewhere. The scalar
9518/// arm is also the test oracle for both fast arms.
9519#[inline]
9520fn q8_row_dot(row: &[u8], act: &SplitAct) -> f32 {
9521    #[cfg(target_arch = "aarch64")]
9522    return row_dot_sdot(row, act);
9523    #[cfg(target_arch = "x86_64")]
9524    return row_dot_avx2(row, act);
9525    #[allow(unreachable_code)]
9526    q8_row_dot_scalar(row, act)
9527}
9528
9529#[allow(dead_code)]
9530fn q8_row_dot_scalar(row: &[u8], act: &SplitAct) -> f32 {
9531    let mut acc = 0i32;
9532    for (k, &b) in row.iter().enumerate() {
9533        acc += (b as i8) as i32 * act.xq[k] as i32;
9534    }
9535    let mut acc = acc as f32 * act.sx;
9536    for &(j, xv) in &act.outliers {
9537        acc += (row[j] as i8) as f32 * xv;
9538    }
9539    acc
9540}
9541
9542/// SDOT row dot with exact outlier correction:
9543/// `dot = sdot(w, xq)·sx + Σ_outl w[j]·x[j]` (then × row_scale by caller).
9544#[cfg(target_arch = "aarch64")]
9545#[inline]
9546fn row_dot_sdot(row: &[u8], act: &SplitAct) -> f32 {
9547    let mut acc = unsafe { dot_i8_sdot(row, &act.xq) } as f32 * act.sx;
9548    for &(j, xv) in &act.outliers {
9549        acc += (row[j] as i8) as f32 * xv;
9550    }
9551    acc
9552}
9553
9554/// One q4 row via SDOT: each 32-group's nibbles unpack to centered i8
9555/// (nib−8 ∈ [−8,7]), int8×int8 `sdot` against the pre-quantized
9556/// activation group, × the group's f16 scale. Returns Σ_g dot_g·s_g;
9557/// the caller multiplies by the activation scale and adds the exact
9558/// outlier terms (port of vmfcore `dot_q4_block_sdot`, +23% measured).
9559/// Nibble order matches the writer: element 2k = low nibble, 2k+1 = high
9560/// → zip(lo,hi) restores flat order.
9561#[cfg(target_arch = "aarch64")]
9562#[target_feature(enable = "neon,dotprod")]
9563unsafe fn dot_q4_row_sdot(packed: &[u8], scales: &[u8], g0: usize, gpr: usize, xq: &[i8]) -> f32 {
9564    // SAFETY: callers uphold slice-length contracts (16 packed bytes and
9565    // 2 scale bytes per group; xq.len() == gpr·GROUP_SIZE).
9566    unsafe {
9567        use core::arch::aarch64::*;
9568        use core::arch::asm;
9569        let lomask = vdupq_n_u8(0x0F);
9570        let eight = vdupq_n_s8(8);
9571        let mut acc = 0f32;
9572        for gi in 0..gpr {
9573            let g = g0 + gi;
9574            let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
9575            let b = vld1q_u8(packed.as_ptr().add(g * 16));
9576            let lo = vandq_u8(b, lomask);
9577            let hi = vshrq_n_u8::<4>(b);
9578            let e0 = vsubq_s8(vreinterpretq_s8_u8(vzip1q_u8(lo, hi)), eight);
9579            let e1 = vsubq_s8(vreinterpretq_s8_u8(vzip2q_u8(lo, hi)), eight);
9580            let x0 = vld1q_s8(xq.as_ptr().add(gi * GROUP_SIZE));
9581            let x1 = vld1q_s8(xq.as_ptr().add(gi * GROUP_SIZE + 16));
9582            let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
9583            asm!(
9584                "sdot {a0:v}.4s, {e0:v}.16b, {x0:v}.16b",
9585                "sdot {a1:v}.4s, {e1:v}.16b, {x1:v}.16b",
9586                a0 = inout(vreg) a0, a1 = inout(vreg) a1,
9587                e0 = in(vreg) e0, x0 = in(vreg) x0, e1 = in(vreg) e1, x1 = in(vreg) x1,
9588                options(pure, nomem, nostack),
9589            );
9590            acc += vaddvq_s32(vaddq_s32(a0, a1)) as f32 * s;
9591        }
9592        acc
9593    }
9594}
9595
9596/// Two-activation q4 row via SDOT: the nibble unpack (the expensive
9597/// part) happens ONCE per group; both pre-quantized activations are
9598/// dotted against the same centered i8 registers. Per-lane math matches
9599/// `dot_q4_row_sdot` exactly.
9600#[cfg(target_arch = "aarch64")]
9601#[target_feature(enable = "neon,dotprod")]
9602unsafe fn dot_q4_row_sdot2(
9603    packed: &[u8],
9604    scales: &[u8],
9605    g0: usize,
9606    gpr: usize,
9607    xq1: &[i8],
9608    xq2: &[i8],
9609) -> (f32, f32) {
9610    // SAFETY: callers uphold slice-length contracts (16 packed bytes and
9611    // 2 scale bytes per group; xq*.len() == gpr·GROUP_SIZE).
9612    unsafe {
9613        use core::arch::aarch64::*;
9614        use core::arch::asm;
9615        let lomask = vdupq_n_u8(0x0F);
9616        let eight = vdupq_n_s8(8);
9617        let (mut acc1, mut acc2) = (0f32, 0f32);
9618        for gi in 0..gpr {
9619            let g = g0 + gi;
9620            let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
9621            let b = vld1q_u8(packed.as_ptr().add(g * 16));
9622            let lo = vandq_u8(b, lomask);
9623            let hi = vshrq_n_u8::<4>(b);
9624            let e0 = vsubq_s8(vreinterpretq_s8_u8(vzip1q_u8(lo, hi)), eight);
9625            let e1 = vsubq_s8(vreinterpretq_s8_u8(vzip2q_u8(lo, hi)), eight);
9626            let x10 = vld1q_s8(xq1.as_ptr().add(gi * GROUP_SIZE));
9627            let x11 = vld1q_s8(xq1.as_ptr().add(gi * GROUP_SIZE + 16));
9628            let x20 = vld1q_s8(xq2.as_ptr().add(gi * GROUP_SIZE));
9629            let x21 = vld1q_s8(xq2.as_ptr().add(gi * GROUP_SIZE + 16));
9630            let (mut a0, mut a1, mut b0, mut b1) = (
9631                vdupq_n_s32(0),
9632                vdupq_n_s32(0),
9633                vdupq_n_s32(0),
9634                vdupq_n_s32(0),
9635            );
9636            asm!(
9637                "sdot {a0:v}.4s, {e0:v}.16b, {x10:v}.16b",
9638                "sdot {a1:v}.4s, {e1:v}.16b, {x11:v}.16b",
9639                "sdot {b0:v}.4s, {e0:v}.16b, {x20:v}.16b",
9640                "sdot {b1:v}.4s, {e1:v}.16b, {x21:v}.16b",
9641                a0 = inout(vreg) a0, a1 = inout(vreg) a1,
9642                b0 = inout(vreg) b0, b1 = inout(vreg) b1,
9643                e0 = in(vreg) e0, e1 = in(vreg) e1,
9644                x10 = in(vreg) x10, x11 = in(vreg) x11,
9645                x20 = in(vreg) x20, x21 = in(vreg) x21,
9646                options(pure, nomem, nostack),
9647            );
9648            acc1 += vaddvq_s32(vaddq_s32(a0, a1)) as f32 * s;
9649            acc2 += vaddvq_s32(vaddq_s32(b0, b1)) as f32 * s;
9650        }
9651        (acc1, acc2)
9652    }
9653}
9654
9655// ───────────────────── fused int8 kernels ─────────────────────
9656
9657/// `acc += w · row` where the row is centered i8 — NEON widen+fma on
9658/// aarch64, scalar elsewhere. The KV-cache q8 value path rides on this.
9659#[inline]
9660pub(crate) fn axpy_i8_f32(acc: &mut [f32], row: &[i8], w: f32) {
9661    #[cfg(target_arch = "aarch64")]
9662    unsafe {
9663        return axpy_i8_f32_neon(acc, row, w);
9664    }
9665    #[cfg(target_arch = "x86_64")]
9666    if avx2_enabled() {
9667        return unsafe { axpy_i8_f32_avx2(acc, row, w) };
9668    }
9669    #[allow(unreachable_code)]
9670    {
9671        for (a, &b) in acc.iter_mut().zip(row) {
9672            *a += w * b as f32;
9673        }
9674    }
9675}
9676
9677/// i8→f32 axpy via AVX2/FMA (x86 mirror of `axpy_i8_f32_neon`).
9678#[cfg(target_arch = "x86_64")]
9679#[target_feature(enable = "avx2,fma")]
9680unsafe fn axpy_i8_f32_avx2(acc: &mut [f32], row: &[i8], w: f32) {
9681    // SAFETY: callers uphold slice-length contracts (see call sites).
9682    unsafe {
9683        use core::arch::x86_64::*;
9684        let n = acc.len().min(row.len());
9685        let ap = acc.as_mut_ptr();
9686        let rp = row.as_ptr();
9687        let wv = _mm256_set1_ps(w);
9688        let mut j = 0usize;
9689        while j + 16 <= n {
9690            let rb = _mm_loadu_si128(rp.add(j) as *const __m128i);
9691            let lo = _mm256_cvtepi32_ps(_mm256_cvtepi8_epi32(rb));
9692            let hi = _mm256_cvtepi32_ps(_mm256_cvtepi8_epi32(_mm_srli_si128::<8>(rb)));
9693            let v0 = _mm256_fmadd_ps(wv, lo, _mm256_loadu_ps(ap.add(j)));
9694            let v1 = _mm256_fmadd_ps(wv, hi, _mm256_loadu_ps(ap.add(j + 8)));
9695            _mm256_storeu_ps(ap.add(j), v0);
9696            _mm256_storeu_ps(ap.add(j + 8), v1);
9697            j += 16;
9698        }
9699        while j < n {
9700            *ap.add(j) += w * (*rp.add(j)) as f32;
9701            j += 1;
9702        }
9703    }
9704}
9705
9706#[cfg(target_arch = "aarch64")]
9707#[target_feature(enable = "neon")]
9708unsafe fn axpy_i8_f32_neon(acc: &mut [f32], row: &[i8], w: f32) {
9709    // SAFETY: callers uphold slice-length contracts (see call sites).
9710    unsafe {
9711        use core::arch::aarch64::*;
9712        let n = acc.len().min(row.len());
9713        let ap = acc.as_mut_ptr();
9714        let rp = row.as_ptr();
9715        let wv = vdupq_n_f32(w);
9716        let mut j = 0usize;
9717        while j + 16 <= n {
9718            let rb = vld1q_s8(rp.add(j));
9719            let lo = vmovl_s8(vget_low_s8(rb));
9720            let hi = vmovl_s8(vget_high_s8(rb));
9721            for (off, half) in [(0, lo), (8, hi)] {
9722                let f0 = vcvtq_f32_s32(vmovl_s16(vget_low_s16(half)));
9723                let f1 = vcvtq_f32_s32(vmovl_s16(vget_high_s16(half)));
9724                let o = j + off;
9725                vst1q_f32(ap.add(o), vfmaq_f32(vld1q_f32(ap.add(o)), wv, f0));
9726                vst1q_f32(ap.add(o + 4), vfmaq_f32(vld1q_f32(ap.add(o + 4)), wv, f1));
9727            }
9728            j += 16;
9729        }
9730        while j < n {
9731            *ap.add(j) += w * (*rp.add(j)) as f32;
9732            j += 1;
9733        }
9734    }
9735}
9736
9737/// i8 row · f32 x. NEON on aarch64 (ported from vmfcore `dot_i8_f32_neon`,
9738/// ≈9× scalar), scalar elsewhere.
9739#[inline]
9740pub(crate) fn dot_i8_f32(w: &[u8], x: &[f32]) -> f32 {
9741    #[cfg(target_arch = "aarch64")]
9742    unsafe {
9743        return dot_i8_f32_neon(w, x);
9744    }
9745    #[cfg(target_arch = "x86_64")]
9746    if avx2_enabled() {
9747        return unsafe { dot_i8_f32_avx2(w, x) };
9748    }
9749    #[allow(unreachable_code)]
9750    {
9751        let mut sum = 0.0f32;
9752        for (j, &b) in w.iter().enumerate() {
9753            sum += (b as i8) as f32 * x[j];
9754        }
9755        sum
9756    }
9757}
9758
9759/// i8 row · (x ⊙ col_field) — the q8_2f row dot with the θ col-field
9760/// folded into the product (no prescaled copy of x). NEON on aarch64,
9761/// scalar elsewhere. Used by the active-neuron path `row_dot`.
9762#[inline]
9763fn dot_i8_col_f32(w: &[u8], x: &[f32], col: &[f32]) -> f32 {
9764    #[cfg(target_arch = "aarch64")]
9765    unsafe {
9766        return dot_i8_col_f32_neon(w, x, col);
9767    }
9768    #[allow(unreachable_code)]
9769    {
9770        let mut sum = 0.0f32;
9771        for (j, &b) in w.iter().enumerate() {
9772            sum += (b as i8) as f32 * x[j] * col[j];
9773        }
9774        sum
9775    }
9776}
9777
9778#[cfg(target_arch = "aarch64")]
9779#[target_feature(enable = "neon")]
9780unsafe fn dot_i8_col_f32_neon(w: &[u8], x: &[f32], col: &[f32]) -> f32 {
9781    // SAFETY: callers uphold slice-length contracts (see call sites).
9782    unsafe {
9783        use core::arch::aarch64::*;
9784        let n = x.len();
9785        let wp = w.as_ptr() as *const i8;
9786        let xp = x.as_ptr();
9787        let cp = col.as_ptr();
9788        let (mut a0, mut a1, mut a2, mut a3) = (
9789            vdupq_n_f32(0.0),
9790            vdupq_n_f32(0.0),
9791            vdupq_n_f32(0.0),
9792            vdupq_n_f32(0.0),
9793        );
9794        let mut j = 0usize;
9795        while j + 16 <= n {
9796            let wb = vld1q_s8(wp.add(j));
9797            let lo = vmovl_s8(vget_low_s8(wb));
9798            let hi = vmovl_s8(vget_high_s8(wb));
9799            let w0 = vcvtq_f32_s32(vmovl_s16(vget_low_s16(lo)));
9800            let w1 = vcvtq_f32_s32(vmovl_s16(vget_high_s16(lo)));
9801            let w2 = vcvtq_f32_s32(vmovl_s16(vget_low_s16(hi)));
9802            let w3 = vcvtq_f32_s32(vmovl_s16(vget_high_s16(hi)));
9803            a0 = vfmaq_f32(
9804                a0,
9805                w0,
9806                vmulq_f32(vld1q_f32(xp.add(j)), vld1q_f32(cp.add(j))),
9807            );
9808            a1 = vfmaq_f32(
9809                a1,
9810                w1,
9811                vmulq_f32(vld1q_f32(xp.add(j + 4)), vld1q_f32(cp.add(j + 4))),
9812            );
9813            a2 = vfmaq_f32(
9814                a2,
9815                w2,
9816                vmulq_f32(vld1q_f32(xp.add(j + 8)), vld1q_f32(cp.add(j + 8))),
9817            );
9818            a3 = vfmaq_f32(
9819                a3,
9820                w3,
9821                vmulq_f32(vld1q_f32(xp.add(j + 12)), vld1q_f32(cp.add(j + 12))),
9822            );
9823            j += 16;
9824        }
9825        let mut sum = vaddvq_f32(vaddq_f32(vaddq_f32(a0, a1), vaddq_f32(a2, a3)));
9826        while j < n {
9827            sum += (*wp.add(j)) as f32 * *xp.add(j) * *cp.add(j);
9828            j += 1;
9829        }
9830        sum
9831    }
9832}
9833
9834#[cfg(target_arch = "aarch64")]
9835#[target_feature(enable = "neon")]
9836unsafe fn dot_i8_f32_neon(w: &[u8], x: &[f32]) -> f32 {
9837    // SAFETY: callers uphold slice-length contracts (see call sites).
9838    unsafe {
9839        use core::arch::aarch64::*;
9840        let n = x.len();
9841        let wp = w.as_ptr() as *const i8;
9842        let xp = x.as_ptr();
9843        let (mut a0, mut a1, mut a2, mut a3) = (
9844            vdupq_n_f32(0.0),
9845            vdupq_n_f32(0.0),
9846            vdupq_n_f32(0.0),
9847            vdupq_n_f32(0.0),
9848        );
9849        let mut j = 0usize;
9850        while j + 16 <= n {
9851            let wb = vld1q_s8(wp.add(j));
9852            let lo = vmovl_s8(vget_low_s8(wb));
9853            let hi = vmovl_s8(vget_high_s8(wb));
9854            let w0 = vcvtq_f32_s32(vmovl_s16(vget_low_s16(lo)));
9855            let w1 = vcvtq_f32_s32(vmovl_s16(vget_high_s16(lo)));
9856            let w2 = vcvtq_f32_s32(vmovl_s16(vget_low_s16(hi)));
9857            let w3 = vcvtq_f32_s32(vmovl_s16(vget_high_s16(hi)));
9858            a0 = vfmaq_f32(a0, w0, vld1q_f32(xp.add(j)));
9859            a1 = vfmaq_f32(a1, w1, vld1q_f32(xp.add(j + 4)));
9860            a2 = vfmaq_f32(a2, w2, vld1q_f32(xp.add(j + 8)));
9861            a3 = vfmaq_f32(a3, w3, vld1q_f32(xp.add(j + 12)));
9862            j += 16;
9863        }
9864        let mut sum = vaddvq_f32(vaddq_f32(vaddq_f32(a0, a1), vaddq_f32(a2, a3)));
9865        while j < n {
9866            sum += (*wp.add(j)) as f32 * *xp.add(j);
9867            j += 1;
9868        }
9869        sum
9870    }
9871}
9872
9873#[allow(clippy::too_many_arguments)]
9874fn qmatvec(
9875    q: &[u8],
9876    rep: &[u8],
9877    row_scale: &[f32],
9878    x: &[f32],
9879    col_field: &[f32],
9880    dtype: TensorDtype,
9881    rows: usize,
9882    cols: usize,
9883    out: &mut [f32],
9884    pool: Option<&Pool>,
9885) {
9886    debug_assert_eq!(out.len(), rows);
9887    #[cfg(not(target_arch = "aarch64"))]
9888    let _ = rep;
9889
9890    #[cfg(target_arch = "aarch64")]
9891    if sdot_enabled() {
9892        let act = if dtype == TensorDtype::Q8_2f {
9893            split_act_q8_2f(x, col_field)
9894        } else {
9895            split_act(x)
9896        };
9897        let out_addr = SendMut(out.as_mut_ptr());
9898        let run_range = |start: usize, end: usize| {
9899            q8_range_sdot(q, rep, row_scale, &act, cols, out_addr, start, end)
9900        };
9901        match pool {
9902            Some(pool) if rows >= 256 => pool.run_rows(rows, &run_range),
9903            _ => run_range(0, rows),
9904        }
9905        return;
9906    }
9907    // x86 A8W8 via AVX2 maddubs — same quantized-activation contract as
9908    // the SDOT path (CMF_AVX2=0 keeps the exact i8×f32 loop).
9909    #[cfg(target_arch = "x86_64")]
9910    if avx2_a8w8_enabled() {
9911        let act = if dtype == TensorDtype::Q8_2f {
9912            split_act_q8_2f(x, col_field)
9913        } else {
9914            split_act(x)
9915        };
9916        let out_addr = SendMut(out.as_mut_ptr());
9917        let run_range = |start: usize, end: usize| {
9918            q8_range_avx2(q, row_scale, &act, cols, out_addr, start, end)
9919        };
9920        match pool {
9921            Some(pool) if rows >= 256 => pool.run_rows(rows, &run_range),
9922            _ => run_range(0, rows),
9923        }
9924        return;
9925    }
9926
9927    prescale_with(x, col_field, dtype, 1, |xs| {
9928        let out_addr = SendMut(out.as_mut_ptr());
9929        let run_range = move |start: usize, end: usize| {
9930            for o in start..end {
9931                let v = dot_i8_f32(&q[o * cols..(o + 1) * cols], xs) * row_scale[o];
9932                // SAFETY: disjoint row ranges per worker.
9933                unsafe { *out_addr.at(o) = v };
9934            }
9935        };
9936        match pool {
9937            Some(pool) if rows >= 256 => pool.run_rows(rows, &run_range),
9938            _ => run_range(0, rows),
9939        }
9940    });
9941}
9942
9943#[allow(clippy::too_many_arguments)]
9944fn qmatvec2(
9945    q: &[u8],
9946    row_scale: &[f32],
9947    x1: &[f32],
9948    x2: &[f32],
9949    col_field: &[f32],
9950    dtype: TensorDtype,
9951    rows: usize,
9952    cols: usize,
9953    o1: &mut [f32],
9954    o2: &mut [f32],
9955    pool: Option<&Pool>,
9956) {
9957    #[cfg(target_arch = "aarch64")]
9958    if sdot_enabled() {
9959        let a1s = if dtype == TensorDtype::Q8_2f {
9960            split_act_q8_2f(x1, col_field)
9961        } else {
9962            split_act(x1)
9963        };
9964        let a2s = if dtype == TensorDtype::Q8_2f {
9965            split_act_q8_2f(x2, col_field)
9966        } else {
9967            split_act(x2)
9968        };
9969        let p1 = SendMut(o1.as_mut_ptr());
9970        let p2 = SendMut(o2.as_mut_ptr());
9971        let run_range = |start: usize, end: usize| {
9972            q8_range2_sdot(q, row_scale, &a1s, &a2s, cols, p1, p2, start, end)
9973        };
9974        match pool {
9975            Some(pool) if rows >= 256 => pool.run_rows(rows, &run_range),
9976            _ => run_range(0, rows),
9977        }
9978        return;
9979    }
9980    #[cfg(target_arch = "x86_64")]
9981    if avx2_a8w8_enabled() {
9982        let a1s = if dtype == TensorDtype::Q8_2f {
9983            split_act_q8_2f(x1, col_field)
9984        } else {
9985            split_act(x1)
9986        };
9987        let a2s = if dtype == TensorDtype::Q8_2f {
9988            split_act_q8_2f(x2, col_field)
9989        } else {
9990            split_act(x2)
9991        };
9992        let p1 = SendMut(o1.as_mut_ptr());
9993        let p2 = SendMut(o2.as_mut_ptr());
9994        let run_range = |start: usize, end: usize| {
9995            q8_range2_avx2(q, row_scale, &a1s, &a2s, cols, p1, p2, start, end)
9996        };
9997        match pool {
9998            Some(pool) if rows >= 256 => pool.run_rows(rows, &run_range),
9999            _ => run_range(0, rows),
10000        }
10001        return;
10002    }
10003
10004    prescale_with(x1, col_field, dtype, 1, |x1s| {
10005        prescale_with(x2, col_field, dtype, 2, |x2s| {
10006            let p1 = SendMut(o1.as_mut_ptr());
10007            let p2 = SendMut(o2.as_mut_ptr());
10008            let run_range = move |start: usize, end: usize| {
10009                for o in start..end {
10010                    let row = &q[o * cols..(o + 1) * cols];
10011                    let s1 = dot_i8_f32(row, x1s) * row_scale[o];
10012                    let s2 = dot_i8_f32(row, x2s) * row_scale[o];
10013                    // SAFETY: disjoint row ranges per worker.
10014                    unsafe {
10015                        *p1.at(o) = s1;
10016                        *p2.at(o) = s2;
10017                    }
10018                }
10019            };
10020            match pool {
10021                Some(pool) if rows >= 256 => pool.run_rows(rows, &run_range),
10022                _ => run_range(0, rows),
10023            }
10024        });
10025    });
10026}
10027
10028#[derive(Clone, Copy)]
10029struct SendMut(*mut f32);
10030unsafe impl Send for SendMut {}
10031unsafe impl Sync for SendMut {}
10032
10033impl SendMut {
10034    #[inline]
10035    fn at(self, i: usize) -> *mut f32 {
10036        unsafe { self.0.add(i) }
10037    }
10038}
10039
10040#[cfg(test)]
10041mod tests {
10042    use super::*;
10043
10044    #[test]
10045    fn q2tp_i8_dot_matches_exact_on_grid() {
10046        // On-grid activations (±1 → sx=1/127, xq=±127 dequantizes
10047        // exactly, no outliers) must make the integer path agree with
10048        // the exact scalar walk to f32 rounding.
10049        let (rows, cols) = (5, 64);
10050        let gpr = cols / GROUP_SIZE;
10051        // Synthetic codes plane + a flat ladder: scales_into is not under
10052        // test here, so drive dot_q2tp_row_i8 / q2tp_row_exact directly
10053        // with hand-made scales.
10054        let chunks: Vec<u8> = (0..rows * gpr * Q2TP_CHUNK)
10055            .map(|i| (i as u32).wrapping_mul(2654435761) as u8)
10056            .collect();
10057        let scales: Vec<f32> = (0..gpr).map(|g| 0.5 + g as f32 * 0.25).collect();
10058        let x: Vec<f32> = (0..cols)
10059            .map(|i| if i % 3 == 0 { -1.0 } else { 1.0 })
10060            .collect();
10061        let act = split_act(&x);
10062        assert!(
10063            act.outliers.is_empty(),
10064            "on-grid input must have no outliers"
10065        );
10066        let gsum = q1_group_sums(&act.xq, gpr);
10067        for r in 0..rows {
10068            let exact = q2tp_row_exact(&chunks, r, gpr, &x, &scales);
10069            let fast = dot_q2tp_row_i8(&chunks, r, gpr, &act.xq, &gsum, &scales) * act.sx;
10070            assert!(
10071                (exact - fast).abs() <= exact.abs() * 1e-5 + 1e-5,
10072                "row {r}: exact {exact} vs i8 {fast}"
10073            );
10074        }
10075    }
10076
10077    #[test]
10078    fn q8_row_dot_fast_matches_scalar() {
10079        // The per-arch fast dot must agree with the exact scalar oracle
10080        // (same contract the fused q8 FFN arm rides on).
10081        let cols = 96;
10082        let row: Vec<u8> = (0..cols)
10083            .map(|i| ((i * 37 % 251) - 125) as i8 as u8)
10084            .collect();
10085        let x: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.13).sin()).collect();
10086        let act = split_act(&x);
10087        let fast = q8_row_dot(&row, &act);
10088        let scalar = q8_row_dot_scalar(&row, &act);
10089        assert!(
10090            (fast - scalar).abs() <= scalar.abs() * 1e-5 + 1e-5,
10091            "fast {fast} vs scalar {scalar}"
10092        );
10093    }
10094
10095    #[test]
10096    fn f32_matvec_matches_matvec_rows_bitexact() {
10097        let (rows, cols) = (300, 40);
10098        let w: Vec<f32> = (0..rows * cols).map(|i| (i as f32 * 0.017).sin()).collect();
10099        let x: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.05).cos()).collect();
10100        let qt = QTensor::from_f32(w.clone(), rows, cols);
10101
10102        let mut a = vec![0.0f32; rows];
10103        matvec_rows(None, &w, &x, &mut a);
10104        let mut b = vec![0.0f32; rows];
10105        qt.matvec(&x, &mut b, None);
10106        assert_eq!(a, b);
10107    }
10108
10109    #[test]
10110    fn sdot_kernel_exact_on_grid() {
10111        // Activations already on the i8 grid (±1 with amax=1 → sx=1/127,
10112        // xq=±127 dequantizes EXACTLY) → the SDOT path must match the
10113        // exact f32 dot to float rounding. This isolates kernel
10114        // correctness from quantization noise.
10115        eprintln!("sdot_enabled = {}", sdot_enabled());
10116        let (rows, cols) = (9, 80); // odd rows → exercises 4-row + tail
10117        let w: Vec<u8> = (0..rows * cols)
10118            .map(|i| (((i * 37) % 251) as i32 - 125) as i8 as u8)
10119            .collect();
10120        let scales: Vec<f32> = (0..rows).map(|o| 0.005 + o as f32 * 0.001).collect();
10121        let x: Vec<f32> = (0..cols)
10122            .map(|i| match i % 3 {
10123                0 => 1.0,
10124                1 => -1.0,
10125                _ => 0.0,
10126            })
10127            .collect();
10128        let mut a = vec![0.0f32; rows];
10129        qmatvec(
10130            &w,
10131            &[],
10132            &scales,
10133            &x,
10134            &[],
10135            TensorDtype::Q8Row,
10136            rows,
10137            cols,
10138            &mut a,
10139            None,
10140        );
10141        for o in 0..rows {
10142            let mut acc = 0.0f32;
10143            for j in 0..cols {
10144                acc += (w[o * cols + j] as i8) as f32 * x[j];
10145            }
10146            let expect = acc * scales[o];
10147            assert!(
10148                (a[o] - expect).abs() < 1e-3 * expect.abs().max(1e-3),
10149                "row {o}: {} vs {expect}",
10150                a[o]
10151            );
10152        }
10153    }
10154
10155    #[test]
10156    fn q1_tbl_fast_path_matches_reference() {
10157        // gpr = 8 exercises the TBL pair-load fast loop, and the LAST
10158        // row's final 4-tile window trips the 4B-overread guard (the
10159        // payload ends exactly at the last tile) — both paths must
10160        // agree with the dequant reference.
10161        let (rows, cols) = (5, 256);
10162        let gpr = cols / GROUP_SIZE;
10163        let mut bytes = Vec::new();
10164        for t in 0..rows * gpr {
10165            let s = 0.007 + (t % 11) as f32 * 0.004;
10166            bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
10167            for j in 0..4 {
10168                bytes.push(((t * 53 + j * 89 + 7) % 249) as u8);
10169            }
10170        }
10171        let x: Vec<f32> = (0..cols)
10172            .map(|i| if (i * 5) % 7 < 3 { 1.0 } else { -1.0 })
10173            .collect();
10174        let mut w = vec![0.0f32; rows * cols];
10175        cortiq_core::quant::dequant_q1(&bytes, &mut w);
10176        let mut got = vec![0.0f32; rows];
10177        q1_matvec(&bytes, &x, rows, cols, &mut got, None);
10178        for o in 0..rows {
10179            let expect: f32 = (0..cols).map(|j| w[o * cols + j] * x[j]).sum();
10180            assert!(
10181                (got[o] - expect).abs() < 1e-3 * expect.abs().max(1e-3),
10182                "row {o}: {} vs {expect}",
10183                got[o]
10184            );
10185        }
10186        // Blocked 1×4 batch (b=5: one quad + remainder) must equal the
10187        // single-matvec path bit-for-bit.
10188        let b = 5usize;
10189        let mut xs_all = Vec::new();
10190        for bi in 0..b {
10191            xs_all.extend(x.iter().map(|v| if bi % 2 == 0 { *v } else { -*v }));
10192        }
10193        let mut mm = vec![0.0f32; b * rows];
10194        q1_matmat(&bytes, &xs_all, b, rows, cols, &mut mm, None);
10195        for bi in 0..b {
10196            let mut single = vec![0.0f32; rows];
10197            q1_matvec(
10198                &bytes,
10199                &xs_all[bi * cols..(bi + 1) * cols],
10200                rows,
10201                cols,
10202                &mut single,
10203                None,
10204            );
10205            assert_eq!(&mm[bi * rows..(bi + 1) * rows], &single[..], "stream {bi}");
10206        }
10207    }
10208
10209    #[test]
10210    fn q1_kernels_match_exact_reference() {
10211        // Synthetic q1 payload: 6-byte tiles [f16 scale][4B bits].
10212        let (rows, cols) = (7, 96);
10213        let gpr = cols / GROUP_SIZE;
10214        let mut bytes = Vec::new();
10215        for t in 0..rows * gpr {
10216            let s = 0.01 + (t % 13) as f32 * 0.003;
10217            bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
10218            for j in 0..4 {
10219                bytes.push(((t * 31 + j * 97) % 251) as u8);
10220            }
10221        }
10222        // On-grid activations (±1, amax 1) → the SDOT path is exact.
10223        let x: Vec<f32> = (0..cols)
10224            .map(|i| if i % 3 == 0 { 1.0 } else { -1.0 })
10225            .collect();
10226        // Reference through the core dequant.
10227        let mut w = vec![0.0f32; rows * cols];
10228        cortiq_core::quant::dequant_q1(&bytes, &mut w);
10229        let mut expect = vec![0.0f32; rows];
10230        for o in 0..rows {
10231            expect[o] = (0..cols).map(|j| w[o * cols + j] * x[j]).sum();
10232        }
10233        let mut got = vec![0.0f32; rows];
10234        q1_matvec(&bytes, &x, rows, cols, &mut got, None);
10235        for o in 0..rows {
10236            assert!(
10237                (got[o] - expect[o]).abs() < 1e-3 * expect[o].abs().max(1e-3),
10238                "row {o}: {} vs {}",
10239                got[o],
10240                expect[o]
10241            );
10242        }
10243        // Pair and batch paths agree with the single path.
10244        let x2: Vec<f32> = x.iter().map(|v| -v).collect();
10245        let (mut a1, mut a2) = (vec![0.0f32; rows], vec![0.0f32; rows]);
10246        q1_matvec2(&bytes, &x, &x2, rows, cols, &mut a1, &mut a2, None);
10247        assert_eq!(a1, got);
10248        let mut xs = x.clone();
10249        xs.extend_from_slice(&x2);
10250        let mut mm = vec![0.0f32; 2 * rows];
10251        q1_matmat(&bytes, &xs, 2, rows, cols, &mut mm, None);
10252        assert_eq!(&mm[..rows], got.as_slice());
10253        assert_eq!(&mm[rows..], a2.as_slice());
10254    }
10255
10256    #[test]
10257    fn repack_is_bit_identical() {
10258        // The interleaved-repack kernel must produce EXACTLY the same
10259        // bits as the mmap-layout kernel: integer accumulation is order-
10260        // exact, the f32 epilogue is identical. Odd rows exercise the
10261        // tail; direct range calls exercise unaligned pool splits.
10262        let (rows, cols) = (267, 96); // 66 groups + 3 tail rows, cols % 16 == 0
10263        let w: Vec<u8> = (0..rows * cols)
10264            .map(|i| (((i * 89) % 253) as i32 - 126) as i8 as u8)
10265            .collect();
10266        let scales: Vec<f32> = (0..rows).map(|o| 0.003 + o as f32 * 0.0007).collect();
10267        let x: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.37).sin() * 2.0).collect();
10268        let rep = q8_repack_layout(&w, rows, cols);
10269        // Group interleave round-trips.
10270        for g in 0..rows / 4 {
10271            for c in 0..cols / 16 {
10272                for lane in 0..4 {
10273                    assert_eq!(
10274                        &rep[g * 4 * cols + c * 64 + lane * 16
10275                            ..g * 4 * cols + c * 64 + lane * 16 + 16],
10276                        &w[(g * 4 + lane) * cols + c * 16..(g * 4 + lane) * cols + c * 16 + 16],
10277                    );
10278                }
10279            }
10280        }
10281        let mut a = vec![0.0f32; rows];
10282        qmatvec(
10283            &w,
10284            &[],
10285            &scales,
10286            &x,
10287            &[],
10288            TensorDtype::Q8Row,
10289            rows,
10290            cols,
10291            &mut a,
10292            None,
10293        );
10294        let mut b = vec![0.0f32; rows];
10295        qmatvec(
10296            &w,
10297            &rep,
10298            &scales,
10299            &x,
10300            &[],
10301            TensorDtype::Q8Row,
10302            rows,
10303            cols,
10304            &mut b,
10305            None,
10306        );
10307        assert_eq!(a, b, "full-range repack output diverged");
10308
10309        #[cfg(target_arch = "aarch64")]
10310        if sdot_enabled() {
10311            // Unaligned range split (pool workers get arbitrary bounds).
10312            let act = split_act(&x);
10313            let mut c1 = vec![0.0f32; rows];
10314            let mut c2 = vec![0.0f32; rows];
10315            q8_range_sdot(
10316                &w,
10317                &[],
10318                &scales,
10319                &act,
10320                cols,
10321                SendMut(c1.as_mut_ptr()),
10322                3,
10323                rows - 2,
10324            );
10325            q8_range_sdot(
10326                &w,
10327                &rep,
10328                &scales,
10329                &act,
10330                cols,
10331                SendMut(c2.as_mut_ptr()),
10332                3,
10333                rows - 2,
10334            );
10335            assert_eq!(c1, c2, "unaligned-range repack output diverged");
10336        }
10337    }
10338
10339    #[test]
10340    fn sdot_a8w8_noise_is_bounded() {
10341        // Off-grid activations: A8 quantization noise must stay small in
10342        // relative L2 over the whole output (realistic accuracy contract;
10343        // vmfcore measured argmax-identical decode on real models).
10344        let (rows, cols) = (16, 512);
10345        let w: Vec<u8> = (0..rows * cols)
10346            .map(|i| (((i * 37) % 251) as i32 - 125) as i8 as u8)
10347            .collect();
10348        let scales = vec![0.01f32; rows];
10349        let x: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.21).sin()).collect();
10350        let mut a = vec![0.0f32; rows];
10351        qmatvec(
10352            &w,
10353            &[],
10354            &scales,
10355            &x,
10356            &[],
10357            TensorDtype::Q8Row,
10358            rows,
10359            cols,
10360            &mut a,
10361            None,
10362        );
10363        let (mut num, mut den) = (0f64, 0f64);
10364        for o in 0..rows {
10365            let mut acc = 0.0f32;
10366            for j in 0..cols {
10367                acc += (w[o * cols + j] as i8) as f32 * x[j];
10368            }
10369            let expect = acc * scales[o];
10370            num += ((a[o] - expect) as f64).powi(2);
10371            den += (expect as f64).powi(2);
10372        }
10373        let rel = (num / den.max(1e-12)).sqrt();
10374        assert!(rel < 0.05, "A8W8 relative L2 error too high: {rel}");
10375    }
10376
10377    #[test]
10378    fn i8_dot_neon_matches_scalar() {
10379        let n = 100;
10380        let w: Vec<u8> = (0..n).map(|i| ((i * 37 + 11) % 251) as u8).collect();
10381        let x: Vec<f32> = (0..n).map(|i| (i as f32 * 0.13).sin()).collect();
10382        let mut scalar = 0.0f32;
10383        for j in 0..n {
10384            scalar += (w[j] as i8) as f32 * x[j];
10385        }
10386        let fast = dot_i8_f32(&w, &x);
10387        assert!((scalar - fast).abs() < 1e-3 * scalar.abs().max(1.0));
10388    }
10389
10390    /// Fused vbit matvec must match full dequant_vbit + dense matvec.
10391    #[test]
10392    fn vbitmatvec_matches_full_dequant() {
10393        let (rows, cols) = (6, 64);
10394        let ng = cols / GROUP_SIZE;
10395        // Hand-craft: bits per row, f16 scales, packed rows.
10396        let bits: Vec<u8> = vec![3, 4, 5, 6, 8, 4];
10397        let mut bytes = bits.clone();
10398        for g in 0..rows * ng {
10399            let s = 0.02 + 0.001 * g as f32;
10400            bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
10401        }
10402        for r in 0..rows {
10403            let b = bits[r] as usize;
10404            let (mut acc, mut nb) = (0u64, 0usize);
10405            let mut rowbytes = Vec::new();
10406            for i in 0..cols {
10407                let v = ((i * 7 + r * 13) % (1 << b)) as u64;
10408                acc = (acc << b) | v;
10409                nb += b;
10410                while nb >= 8 {
10411                    nb -= 8;
10412                    rowbytes.push(((acc >> nb) & 0xFF) as u8);
10413                }
10414            }
10415            if nb > 0 {
10416                rowbytes.push(((acc << (8 - nb)) & 0xFF) as u8);
10417            }
10418            bytes.extend_from_slice(&rowbytes);
10419        }
10420        let x: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.19).sin()).collect();
10421
10422        let mut reference = vec![0f32; rows * cols];
10423        cortiq_core::quant::dequant_vbit(&bytes, rows, cols, &mut reference).unwrap();
10424        let mut expect = vec![0f32; rows];
10425        for r in 0..rows {
10426            expect[r] = reference[r * cols..(r + 1) * cols]
10427                .iter()
10428                .zip(&x)
10429                .map(|(w, xv)| w * xv)
10430                .sum();
10431        }
10432        let mut got = vec![0f32; rows];
10433        let offsets = vbit_row_offsets(&bytes, rows, cols);
10434        vbitmatvec(&bytes, &offsets, &x, rows, cols, &mut got, None);
10435        // SDOT path quantizes activations to i8 (A8W8): bounded noise,
10436        // same contract as q8 (exact path is pinned by CMF_SDOT=0 in
10437        // the golden-parity gate).
10438        let tol = if a8w8_enabled() { 6e-2 } else { 1e-4 };
10439        let scale = expect.iter().fold(0f32, |m, v| m.max(v.abs())).max(1e-6);
10440        for r in 0..rows {
10441            assert!(
10442                (got[r] - expect[r]).abs() < tol * scale,
10443                "row {r}: {} vs {}",
10444                got[r],
10445                expect[r]
10446            );
10447        }
10448    }
10449
10450    /// Fused q4 matvec must match the reference full-dequant + dense
10451    /// matvec bit-for-bit in structure (same f32 math, group order).
10452    /// vbit matmat: the blocked 1×4 leg must match the per-row path
10453    /// (paired env toggle; larger shape so both code paths engage).
10454    #[test]
10455    #[cfg(target_arch = "x86_64")]
10456    fn vbit_matmat_blocked_matches_per_row() {
10457        let (rows, cols, b) = (64usize, 128usize, 9usize);
10458        let ng = cols / GROUP_SIZE;
10459        let bits: Vec<u8> = (0..rows).map(|r| [3u8, 4, 5, 6][r % 4]).collect();
10460        let mut bytes = bits.clone();
10461        for g in 0..rows * ng {
10462            let sc = 0.02 + 0.0005 * g as f32;
10463            bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(sc).to_le_bytes());
10464        }
10465        for r in 0..rows {
10466            let bw = bits[r] as usize;
10467            let (mut acc, mut nb) = (0u64, 0usize);
10468            let mut rowbytes = Vec::new();
10469            for i in 0..cols {
10470                let v = ((i * 7 + r * 13) % (1 << bw)) as u64;
10471                acc = (acc << bw) | v;
10472                nb += bw;
10473                while nb >= 8 {
10474                    nb -= 8;
10475                    rowbytes.push(((acc >> nb) & 0xFF) as u8);
10476                }
10477            }
10478            if nb > 0 {
10479                rowbytes.push(((acc << (8 - nb)) & 0xFF) as u8);
10480            }
10481            bytes.extend_from_slice(&rowbytes);
10482        }
10483        let x: Vec<f32> = (0..b * cols)
10484            .map(|i| ((i * 13 + 7) % 97) as f32 / 97.0 - 0.5)
10485            .collect();
10486        let offsets = vbit_row_offsets(&bytes, rows, cols);
10487        let mut y_a = vec![0f32; b * rows];
10488        let mut y_b = vec![0f32; b * rows];
10489        unsafe { std::env::set_var("CMF_X86_BLOCKED", "1") };
10490        vbitmatmat(&bytes, &offsets, &x, b, rows, cols, &mut y_a, None);
10491        unsafe { std::env::set_var("CMF_X86_BLOCKED", "0") };
10492        vbitmatmat(&bytes, &offsets, &x, b, rows, cols, &mut y_b, None);
10493        unsafe { std::env::remove_var("CMF_X86_BLOCKED") };
10494        let max_d = y_a
10495            .iter()
10496            .zip(&y_b)
10497            .map(|(p, q)| (p - q).abs())
10498            .fold(0.0f32, f32::max);
10499        assert!(max_d < 1e-4, "vbit blocked ≠ per-row: max|Δ| = {max_d}");
10500    }
10501
10502    /// q4t blocked 1×4 (SDOT on ARM, AVX2 on x86) must equal the
10503    /// per-row path exactly: same nibble unpack, same group order,
10504    /// same f32 accumulation — batch == matvec bit-for-bit. b=9 covers
10505    /// two full 1×4 blocks plus a remainder through the single-row
10506    /// kernel. (Both paths produce identical output, so the shared
10507    /// CMF_X86_BLOCKED env var racing with other tests cannot flip
10508    /// the verdict — worst case both sides take the same path.)
10509    #[test]
10510    fn q4t_matmat_blocked_matches_per_row() {
10511        let (rows, cols, b) = (16usize, 64usize, 9usize);
10512        let gpr = cols / GROUP_SIZE;
10513        let mut bytes = vec![0u8; rows * gpr * Q4_TILE];
10514        for r in 0..rows {
10515            for g in 0..gpr {
10516                let t = (r * gpr + g) * Q4_TILE;
10517                let sc = 0.02 + 0.001 * (r * gpr + g) as f32;
10518                bytes[t..t + 2].copy_from_slice(&cortiq_core::quant::f32_to_f16(sc).to_le_bytes());
10519                for k in 0..16 {
10520                    bytes[t + 2 + k] = ((r * 31 + g * 7 + k * 13) % 251) as u8;
10521                }
10522            }
10523        }
10524        let x: Vec<f32> = (0..b * cols)
10525            .map(|i| ((i * 13 + 7) % 97) as f32 / 97.0 - 0.5)
10526            .collect();
10527        let mut y_blk = vec![0f32; b * rows];
10528        let mut y_row = vec![0f32; b * rows];
10529        unsafe { std::env::set_var("CMF_X86_BLOCKED", "1") };
10530        q4t_matmat(&bytes, &x, b, rows, cols, &mut y_blk, None);
10531        unsafe { std::env::set_var("CMF_X86_BLOCKED", "0") };
10532        q4t_matmat(&bytes, &x, b, rows, cols, &mut y_row, None);
10533        unsafe { std::env::remove_var("CMF_X86_BLOCKED") };
10534        assert_eq!(y_blk, y_row, "q4t blocked 1x4 ≠ per-row");
10535    }
10536
10537    /// The wide-batch Accelerate arm of q4t_matmat vs a brute-force
10538    /// f32 dequant matmul: both are f32 GEMMs, so only reduction
10539    /// order differs — tight tolerance.
10540    /// A synthetic q4tp payload: random nibbles plus a per-row ladder whose
10541    /// span varies row to row, so the codes actually exercise the full 0..31
10542    /// range rather than clustering on one rung.
10543    fn synth_q4tp(rows: usize, cols: usize) -> Vec<u8> {
10544        use cortiq_core::quant::{f32_to_f16, q4tp_code_stride, q4tp_put_code};
10545        let gpr = cols / GROUP_SIZE;
10546        let stride = q4tp_code_stride(gpr);
10547        let (params_off, codes_off, _) = q4tp_sections(rows, cols);
10548        let mut b = vec![0u8; codes_off + rows * stride];
10549        for r in 0..rows {
10550            for g in 0..gpr {
10551                let t = (r * gpr + g) * Q4TP_NIB;
10552                for k in 0..16 {
10553                    b[t + k] = ((r * 31 + g * 7 + k * 13) % 251) as u8;
10554                }
10555            }
10556            let lo = -6.0 - 0.03 * (r % 17) as f32;
10557            let step = 0.01 + 0.004 * (r % 11) as f32;
10558            let p = params_off + r * 4;
10559            b[p..p + 2].copy_from_slice(&f32_to_f16(lo).to_le_bytes());
10560            b[p + 2..p + 4].copy_from_slice(&f32_to_f16(step).to_le_bytes());
10561            let crow = &mut b[codes_off + r * stride..codes_off + (r + 1) * stride];
10562            for g in 0..gpr {
10563                q4tp_put_code(crow, g, (r * 5 + g * 3) % 32);
10564            }
10565        }
10566        b
10567    }
10568
10569    /// The same weights re-expressed as q4_tiled, so the proven kernel can
10570    /// be the reference: each tile stores the ladder scale its code selects.
10571    /// Only the f16 rounding of that scale separates the two payloads.
10572    fn q4tp_as_q4t(bytes: &[u8], rows: usize, cols: usize) -> Vec<u8> {
10573        let gpr = cols / GROUP_SIZE;
10574        let v = Q4tpView::new(bytes, rows, cols);
10575        let mut out = vec![0u8; rows * gpr * Q4_TILE];
10576        let mut sc = vec![0f32; gpr];
10577        for r in 0..rows {
10578            v.scales_into(r, gpr, &mut sc);
10579            for g in 0..gpr {
10580                let t = (r * gpr + g) * Q4_TILE;
10581                let s = sc[g];
10582                out[t..t + 2].copy_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
10583                let src = (r * gpr + g) * Q4TP_NIB;
10584                out[t + 2..t + Q4_TILE].copy_from_slice(&v.nib[src..src + Q4TP_NIB]);
10585            }
10586        }
10587        out
10588    }
10589
10590    /// The exact (`CMF_SDOT=0`) path must reproduce `dequant_q4tp` to f32
10591    /// rounding — that scalar routine is the format's definition, and the
10592    /// kernels re-derive the scale from the ladder independently. Call the
10593    /// row kernel directly: `matmat` picks the int8 arm when a8w8 is on,
10594    /// so routing through it would test the other path by accident.
10595    #[test]
10596    fn q4tp_exact_path_matches_dequant_reference() {
10597        let (rows, cols) = (256usize, 512usize);
10598        let gpr = cols / GROUP_SIZE;
10599        let bytes = synth_q4tp(rows, cols);
10600        let mut w = vec![0f32; rows * cols];
10601        cortiq_core::quant::dequant_q4tp(&bytes, rows, cols, &mut w);
10602
10603        let x: Vec<f32> = (0..cols)
10604            .map(|i| ((i * 13 + 7) % 97) as f32 / 97.0 - 0.5)
10605            .collect();
10606        let v = Q4tpView::new(&bytes, rows, cols);
10607        let mut sc = vec![0f32; gpr];
10608        for r in 0..rows {
10609            v.scales_into(r, gpr, &mut sc);
10610            let got = q4tp_row_exact(v.nib, r, gpr, &x, &sc);
10611            let want: f32 = (0..cols).map(|c| w[r * cols + c] * x[c]).sum();
10612            // These dot products cancel down to ~1e-3 from terms of ~5e-2, so
10613            // the meaningful yardstick is the summed magnitude, not the result:
10614            // against the result any reordering of a 512-term f32 sum "fails".
10615            let mag: f32 = (0..cols).map(|c| (w[r * cols + c] * x[c]).abs()).sum();
10616            assert!(
10617                (got - want).abs() <= 1e-5 * mag,
10618                "row {r}: kernel {got} vs dequant {want}"
10619            );
10620        }
10621    }
10622
10623    /// The int8 (a8w8) path can't be checked against an f32 reference — the
10624    /// activation quantization dominates. Check it against the q4t kernel it
10625    /// was ported from instead, on payloads holding the same weights: that
10626    /// isolates exactly what the port could break (16 B stride, ladder
10627    /// lookup, nibble unpack) from what it deliberately shares.
10628    #[test]
10629    fn q4tp_matvec_matches_the_q4t_kernel_it_was_ported_from() {
10630        let (rows, cols) = (256usize, 512usize);
10631        let bytes = synth_q4tp(rows, cols);
10632        let twin = q4tp_as_q4t(&bytes, rows, cols);
10633        let x: Vec<f32> = (0..cols)
10634            .map(|i| ((i * 13 + 7) % 97) as f32 / 97.0 - 0.5)
10635            .collect();
10636
10637        let mut got = vec![0f32; rows];
10638        q4tp_matvec(&bytes, &x, rows, cols, &mut got, None);
10639        let mut want = vec![0f32; rows];
10640        q4t_matvec(&twin, &x, rows, cols, &mut want, None);
10641
10642        // Scale is f16 in the twin and f32 here, so allow that rounding on
10643        // top of the summed magnitude (same cancellation argument as above).
10644        let mut w = vec![0f32; rows * cols];
10645        cortiq_core::quant::dequant_q4tp(&bytes, rows, cols, &mut w);
10646        for r in 0..rows {
10647            let mag: f32 = (0..cols).map(|c| (w[r * cols + c] * x[c]).abs()).sum();
10648            assert!(
10649                (got[r] - want[r]).abs() <= 1e-3 * mag,
10650                "row {r}: q4tp {} vs q4t {}",
10651                got[r],
10652                want[r]
10653            );
10654        }
10655    }
10656
10657    /// `matmat` carries three arms (Accelerate, blocked int8 1x4, scalar).
10658    /// Batch 5 crosses the blocked kernel's stride, so this exercises the
10659    /// 1x4 path AND its scalar tail in one run — the blocked kernel is new
10660    /// code and its four accumulators are exactly what tends to go wrong.
10661    #[test]
10662    fn q4tp_matmat_matches_the_q4t_kernel_it_was_ported_from() {
10663        let (rows, cols, b) = (256usize, 512usize, 5usize);
10664        let bytes = synth_q4tp(rows, cols);
10665        let twin = q4tp_as_q4t(&bytes, rows, cols);
10666        let xs: Vec<f32> = (0..b * cols)
10667            .map(|i| ((i * 29 + 11) % 89) as f32 / 89.0 - 0.5)
10668            .collect();
10669
10670        let mut got = vec![0f32; b * rows];
10671        q4tp_matmat(&bytes, &xs, b, rows, cols, &mut got, None);
10672        let mut want = vec![0f32; b * rows];
10673        q4t_matmat(&twin, &xs, b, rows, cols, &mut want, None);
10674
10675        let mut w = vec![0f32; rows * cols];
10676        cortiq_core::quant::dequant_q4tp(&bytes, rows, cols, &mut w);
10677        for t in 0..b {
10678            for r in 0..rows {
10679                let mag: f32 = (0..cols)
10680                    .map(|c| (w[r * cols + c] * xs[t * cols + c]).abs())
10681                    .sum();
10682                let (g, wa) = (got[t * rows + r], want[t * rows + r]);
10683                assert!(
10684                    (g - wa).abs() <= 1e-3 * mag,
10685                    "batch {t} row {r}: q4tp {g} vs q4t {wa}"
10686                );
10687            }
10688        }
10689    }
10690
10691    #[test]
10692    fn q4tp_matvec2_matches_the_single_stream_kernel() {
10693        let (rows, cols) = (128usize, 256usize);
10694        let gpr = cols / GROUP_SIZE;
10695        let bytes = synth_q4tp(rows, cols);
10696        let xs: Vec<f32> = (0..2 * cols)
10697            .map(|i| ((i * 29 + 11) % 89) as f32 / 89.0 - 0.5)
10698            .collect();
10699
10700        let (mut o1, mut o2) = (vec![0f32; rows], vec![0f32; rows]);
10701        q4tp_matvec2(
10702            &bytes,
10703            &xs[..cols],
10704            &xs[cols..],
10705            rows,
10706            cols,
10707            &mut o1,
10708            &mut o2,
10709            None,
10710        );
10711
10712        // matvec2 takes the exact path for both streams, so the single-row
10713        // kernel is an exact reference — no tolerance for path differences.
10714        let v = Q4tpView::new(&bytes, rows, cols);
10715        let mut sc = vec![0f32; gpr];
10716        for r in 0..rows {
10717            v.scales_into(r, gpr, &mut sc);
10718            assert_eq!(o1[r], q4tp_row_exact(v.nib, r, gpr, &xs[..cols], &sc));
10719            assert_eq!(o2[r], q4tp_row_exact(v.nib, r, gpr, &xs[cols..], &sc));
10720        }
10721    }
10722
10723    /// q4tp must not COST speed — it exists to save bytes, and a format that
10724    /// trades 7% of a file for a slower model is a bad trade. This guard is
10725    /// here because correctness tests happily passed while `q4tp_matmat` was
10726    /// missing its int8 and Accelerate arms and the model ran 5x slower.
10727    /// Measured on M-series: 0.97-1.04x, i.e. parity (16 B tiles are better
10728    /// aligned than q4t's 18 B, which pays for the scale indirection).
10729    #[test]
10730    fn q4tp_matvec_keeps_pace_with_q4t() {
10731        let (rows, cols) = (4096usize, 3072usize);
10732        let bytes = synth_q4tp(rows, cols);
10733        let twin = q4tp_as_q4t(&bytes, rows, cols);
10734        let x: Vec<f32> = (0..cols).map(|i| (i % 97) as f32 / 97.0 - 0.5).collect();
10735        let mut o = vec![0f32; rows];
10736        let n = 12;
10737        let mut best = (f64::MAX, f64::MAX);
10738        // Interleaved A/B, minimum statistic: this machine throttles, and a
10739        // mean over a thermal ramp reliably indicts whichever ran second.
10740        for _ in 0..3 {
10741            let t0 = std::time::Instant::now();
10742            for _ in 0..n {
10743                q4t_matvec(&twin, &x, rows, cols, &mut o, None);
10744            }
10745            best.0 = best.0.min(t0.elapsed().as_secs_f64());
10746            let t0 = std::time::Instant::now();
10747            for _ in 0..n {
10748                q4tp_matvec(&bytes, &x, rows, cols, &mut o, None);
10749            }
10750            best.1 = best.1.min(t0.elapsed().as_secs_f64());
10751        }
10752        let ratio = best.1 / best.0;
10753        println!(
10754            "q4t {:.3} ms | q4tp {:.3} ms | {ratio:.2}x",
10755            best.0 * 1e3 / n as f64,
10756            best.1 * 1e3 / n as f64
10757        );
10758        assert!(ratio < 2.0, "q4tp matvec {ratio:.2}x slower than q4t");
10759    }
10760
10761    #[cfg(target_os = "macos")]
10762    #[test]
10763    fn q4t_matmat_accel_matches_dequant_reference() {
10764        if !accel_gemm_enabled() {
10765            return; // CMF_ACCEL=0
10766        }
10767        let (rows, cols, b) = (512usize, 1024usize, 8usize); // ≥500K → accel arm
10768        let gpr = cols / GROUP_SIZE;
10769        let mut bytes = vec![0u8; rows * gpr * Q4_TILE];
10770        for r in 0..rows {
10771            for g in 0..gpr {
10772                let t = (r * gpr + g) * Q4_TILE;
10773                let sc = 0.02 + 0.0005 * ((r * gpr + g) % 64) as f32;
10774                bytes[t..t + 2].copy_from_slice(&cortiq_core::quant::f32_to_f16(sc).to_le_bytes());
10775                for k in 0..16 {
10776                    bytes[t + 2 + k] = ((r * 31 + g * 7 + k * 13) % 251) as u8;
10777                }
10778            }
10779        }
10780        let x: Vec<f32> = (0..b * cols)
10781            .map(|i| ((i * 13 + 7) % 97) as f32 / 97.0 - 0.5)
10782            .collect();
10783        let mut got = vec![0f32; b * rows];
10784        q4t_matmat(&bytes, &x, b, rows, cols, &mut got, None);
10785        // Brute-force reference off the same tiles.
10786        let mut w = vec![0f32; rows * cols];
10787        for r in 0..rows {
10788            for g in 0..gpr {
10789                let t = (r * gpr + g) * Q4_TILE;
10790                let s = f16_to_f32(u16::from_le_bytes([bytes[t], bytes[t + 1]]));
10791                for (k, &bb) in bytes[t + 2..t + Q4_TILE].iter().enumerate() {
10792                    w[r * cols + g * GROUP_SIZE + k * 2] = ((bb & 0x0F) as f32 - 8.0) * s;
10793                    w[r * cols + g * GROUP_SIZE + k * 2 + 1] =
10794                        (((bb >> 4) & 0x0F) as f32 - 8.0) * s;
10795                }
10796            }
10797        }
10798        for bi in 0..b {
10799            for r in 0..rows {
10800                let want: f32 = (0..cols).map(|j| x[bi * cols + j] * w[r * cols + j]).sum();
10801                let d = (got[bi * rows + r] - want).abs();
10802                assert!(
10803                    d <= want.abs().max(1.0) * 1e-4,
10804                    "accel q4t GEMM diverged at ({bi},{r}): {} vs {want}",
10805                    got[bi * rows + r]
10806                );
10807            }
10808        }
10809    }
10810
10811    #[test]
10812    fn q4matvec_matches_full_dequant() {
10813        let (rows, cols) = (8, 64);
10814        let groups = rows * cols / GROUP_SIZE;
10815        // Hand-craft a q4_block blob: nibbles then f16 scales.
10816        let mut bytes = Vec::with_capacity(groups * 16 + groups * 2);
10817        for i in 0..groups * 16 {
10818            bytes.push((((i * 7 + 3) % 256) & 0xFF) as u8);
10819        }
10820        for g in 0..groups {
10821            let s = 0.01 + 0.003 * g as f32;
10822            bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
10823        }
10824        let x: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.17).sin()).collect();
10825
10826        let mut reference = vec![0.0f32; rows * cols];
10827        cortiq_core::quant::dequant_q4_block(&bytes, &mut reference);
10828        let mut expect = vec![0.0f32; rows];
10829        for r in 0..rows {
10830            expect[r] = reference[r * cols..(r + 1) * cols]
10831                .iter()
10832                .zip(&x)
10833                .map(|(w, xv)| w * xv)
10834                .sum();
10835        }
10836
10837        let mut got = vec![0.0f32; rows];
10838        q4matvec(&bytes, &x, rows, cols, &mut got, None);
10839        // SDOT path quantizes activations to i8 (A8W8): bounded noise,
10840        // same contract as q8/vbit (exact path is pinned by CMF_SDOT=0
10841        // in the golden-parity gate).
10842        let tol = if a8w8_enabled() { 6e-2 } else { 1e-4 };
10843        let scale = expect.iter().fold(0f32, |m, v| m.max(v.abs())).max(1.0);
10844        for r in 0..rows {
10845            assert!(
10846                (got[r] - expect[r]).abs() < tol * scale,
10847                "row {r}: {} vs {}",
10848                got[r],
10849                expect[r]
10850            );
10851        }
10852    }
10853
10854    /// Fused two-input vbit matvec must equal two single matvecs exactly
10855    /// (same per-lane accumulation order on both scalar and SDOT paths).
10856    #[test]
10857    fn vbitmatvec2_equals_two_singles() {
10858        let (rows, cols) = (6, 64);
10859        let ng = cols / GROUP_SIZE;
10860        let bits: Vec<u8> = vec![3, 4, 5, 6, 8, 4];
10861        let mut bytes = bits.clone();
10862        for g in 0..rows * ng {
10863            let s = 0.02 + 0.001 * g as f32;
10864            bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
10865        }
10866        for r in 0..rows {
10867            let b = bits[r] as usize;
10868            let (mut acc, mut nb) = (0u64, 0usize);
10869            let mut rowbytes = Vec::new();
10870            for i in 0..cols {
10871                let v = ((i * 7 + r * 13) % (1 << b)) as u64;
10872                acc = (acc << b) | v;
10873                nb += b;
10874                while nb >= 8 {
10875                    nb -= 8;
10876                    rowbytes.push(((acc >> nb) & 0xFF) as u8);
10877                }
10878            }
10879            if nb > 0 {
10880                rowbytes.push(((acc << (8 - nb)) & 0xFF) as u8);
10881            }
10882            bytes.extend_from_slice(&rowbytes);
10883        }
10884        let x1: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.19).sin()).collect();
10885        let x2: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.11).cos()).collect();
10886        let offsets = vbit_row_offsets(&bytes, rows, cols);
10887
10888        let (mut a1, mut a2) = (vec![0f32; rows], vec![0f32; rows]);
10889        vbitmatvec(&bytes, &offsets, &x1, rows, cols, &mut a1, None);
10890        vbitmatvec(&bytes, &offsets, &x2, rows, cols, &mut a2, None);
10891        let (mut b1, mut b2) = (vec![0f32; rows], vec![0f32; rows]);
10892        vbitmatvec2(
10893            &bytes, &offsets, &x1, &x2, rows, cols, &mut b1, &mut b2, None,
10894        );
10895        assert_eq!(a1, b1, "fused vbit lane 1 must be bit-identical");
10896        assert_eq!(a2, b2, "fused vbit lane 2 must be bit-identical");
10897    }
10898
10899    /// Fused two-input q4 matvec must equal two single matvecs exactly.
10900    #[test]
10901    fn q4matvec2_equals_two_singles() {
10902        let (rows, cols) = (8, 128);
10903        let groups = rows * cols / GROUP_SIZE;
10904        let mut bytes = Vec::with_capacity(groups * 16 + groups * 2);
10905        for i in 0..groups * 16 {
10906            bytes.push((((i * 7 + 3) % 256) & 0xFF) as u8);
10907        }
10908        for g in 0..groups {
10909            let s = 0.01 + 0.003 * g as f32;
10910            bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
10911        }
10912        // Include an outlier channel so the SDOT correction path is
10913        // exercised in the pair kernel too.
10914        let mut x1: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.17).sin()).collect();
10915        x1[9] = 250.0;
10916        let x2: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.23).cos()).collect();
10917
10918        let (mut a1, mut a2) = (vec![0f32; rows], vec![0f32; rows]);
10919        q4matvec(&bytes, &x1, rows, cols, &mut a1, None);
10920        q4matvec(&bytes, &x2, rows, cols, &mut a2, None);
10921        let (mut b1, mut b2) = (vec![0f32; rows], vec![0f32; rows]);
10922        q4matvec2(&bytes, &x1, &x2, rows, cols, &mut b1, &mut b2, None);
10923        assert_eq!(a1, b1, "fused q4 lane 1 must be bit-identical");
10924        assert_eq!(a2, b2, "fused q4 lane 2 must be bit-identical");
10925    }
10926
10927    /// Multi-matrix job must equal separate matvecs exactly — same
10928    /// kernels, only the dispatch is fused.
10929    #[test]
10930    fn matvec_many_equals_separate_matvecs() {
10931        use crate::pool::Pool;
10932        let (r1, r2, cols) = (300, 200, 64);
10933        let mk = |salt: usize, rows: usize| {
10934            QTensor::from_f32(
10935                (0..rows * cols)
10936                    .map(|i| ((i * 7 + salt) % 97) as f32 / 97.0 - 0.5)
10937                    .collect(),
10938                rows,
10939                cols,
10940            )
10941        };
10942        let (a, b) = (mk(1, r1), mk(5, r2));
10943        let x: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.11).sin()).collect();
10944        let pool = Pool::new(3);
10945
10946        let (mut ea, mut eb) = (vec![0f32; r1], vec![0f32; r2]);
10947        a.matvec(&x, &mut ea, Some(&pool));
10948        b.matvec(&x, &mut eb, Some(&pool));
10949        let (mut ga, mut gb) = (vec![0f32; r1], vec![0f32; r2]);
10950        QTensor::matvec_many([&a, &b], &x, [&mut ga, &mut gb], Some(&pool));
10951        assert_eq!(ea, ga, "fused multi-matrix lane 1 must be bit-identical");
10952        assert_eq!(eb, gb, "fused multi-matrix lane 2 must be bit-identical");
10953    }
10954
10955    /// Batched q4/vbit matmat must equal per-position matvec calls
10956    /// exactly (the fallback it replaced) — same kernels, same order.
10957    #[test]
10958    fn batched_matmat_equals_per_position_matvec() {
10959        let (rows, cols, b) = (8, 64, 5);
10960        // q4 blob.
10961        let groups = rows * cols / GROUP_SIZE;
10962        let mut q4 = Vec::new();
10963        for i in 0..groups * 16 {
10964            q4.push((((i * 7 + 3) % 256) & 0xFF) as u8);
10965        }
10966        for g in 0..groups {
10967            q4.extend_from_slice(
10968                &cortiq_core::quant::f32_to_f16(0.01 + 0.003 * g as f32).to_le_bytes(),
10969            );
10970        }
10971        // vbit blob (mixed widths incl. 8).
10972        let ng = cols / GROUP_SIZE;
10973        let bits: Vec<u8> = vec![3, 4, 5, 6, 8, 4, 5, 3];
10974        let mut vb = bits.clone();
10975        for g in 0..rows * ng {
10976            vb.extend_from_slice(
10977                &cortiq_core::quant::f32_to_f16(0.02 + 0.001 * g as f32).to_le_bytes(),
10978            );
10979        }
10980        for r in 0..rows {
10981            let bw = bits[r] as usize;
10982            let (mut acc, mut nb) = (0u64, 0usize);
10983            let mut rowbytes = Vec::new();
10984            for i in 0..cols {
10985                let v = ((i * 7 + r * 13) % (1 << bw)) as u64;
10986                acc = (acc << bw) | v;
10987                nb += bw;
10988                while nb >= 8 {
10989                    nb -= 8;
10990                    rowbytes.push(((acc >> nb) & 0xFF) as u8);
10991                }
10992            }
10993            if nb > 0 {
10994                rowbytes.push(((acc << (8 - nb)) & 0xFF) as u8);
10995            }
10996            vb.extend_from_slice(&rowbytes);
10997        }
10998        let offsets = vbit_row_offsets(&vb, rows, cols);
10999
11000        let xs: Vec<f32> = (0..b * cols).map(|i| (i as f32 * 0.13).sin()).collect();
11001
11002        // q4: batch vs singles.
11003        let mut got = vec![0f32; b * rows];
11004        q4matmat(&q4, &xs, b, rows, cols, &mut got, None);
11005        for bi in 0..b {
11006            let mut expect = vec![0f32; rows];
11007            q4matvec(
11008                &q4,
11009                &xs[bi * cols..(bi + 1) * cols],
11010                rows,
11011                cols,
11012                &mut expect,
11013                None,
11014            );
11015            assert_eq!(
11016                &got[bi * rows..(bi + 1) * rows],
11017                &expect[..],
11018                "q4 batch pos {bi}"
11019            );
11020        }
11021
11022        // vbit: batch vs singles.
11023        let mut got = vec![0f32; b * rows];
11024        vbitmatmat(&vb, &offsets, &xs, b, rows, cols, &mut got, None);
11025        for bi in 0..b {
11026            let mut expect = vec![0f32; rows];
11027            vbitmatvec(
11028                &vb,
11029                &offsets,
11030                &xs[bi * cols..(bi + 1) * cols],
11031                rows,
11032                cols,
11033                &mut expect,
11034                None,
11035            );
11036            assert_eq!(
11037                &got[bi * rows..(bi + 1) * rows],
11038                &expect[..],
11039                "vbit batch pos {bi}"
11040            );
11041        }
11042    }
11043
11044    /// q4_tiled kernels must produce BIT-identical outputs to the q4
11045    /// split kernels on the same values (same ints, same order — only
11046    /// the byte placement differs).
11047    #[test]
11048    fn q4_tiled_matches_q4_block_bitexact() {
11049        let (rows, cols, b) = (8usize, 128usize, 3usize);
11050        let groups = rows * cols / GROUP_SIZE;
11051        let mut split = Vec::with_capacity(groups * 18);
11052        for i in 0..groups * 16 {
11053            split.push((((i * 7 + 3) % 256) & 0xFF) as u8);
11054        }
11055        for g in 0..groups {
11056            split.extend_from_slice(
11057                &cortiq_core::quant::f32_to_f16(0.01 + 0.003 * g as f32).to_le_bytes(),
11058            );
11059        }
11060        // Re-tile: [scale][nibbles] per group.
11061        let (packed, scales) = split.split_at(groups * 16);
11062        let mut tiled = Vec::with_capacity(groups * Q4_TILE);
11063        for g in 0..groups {
11064            tiled.extend_from_slice(&scales[g * 2..g * 2 + 2]);
11065            tiled.extend_from_slice(&packed[g * 16..(g + 1) * 16]);
11066        }
11067
11068        let mut x1: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.17).sin()).collect();
11069        x1[9] = 250.0; // exercise the outlier path
11070        let x2: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.23).cos()).collect();
11071
11072        let (mut a, mut t) = (vec![0f32; rows], vec![0f32; rows]);
11073        q4matvec(&split, &x1, rows, cols, &mut a, None);
11074        q4t_matvec(&tiled, &x1, rows, cols, &mut t, None);
11075        assert_eq!(a, t, "q4t matvec must match q4 bit-for-bit");
11076
11077        let (mut a1, mut a2) = (vec![0f32; rows], vec![0f32; rows]);
11078        let (mut t1, mut t2) = (vec![0f32; rows], vec![0f32; rows]);
11079        q4matvec2(&split, &x1, &x2, rows, cols, &mut a1, &mut a2, None);
11080        q4t_matvec2(&tiled, &x1, &x2, rows, cols, &mut t1, &mut t2, None);
11081        assert_eq!(a1, t1);
11082        assert_eq!(a2, t2);
11083
11084        let xs: Vec<f32> = (0..b * cols).map(|i| (i as f32 * 0.13).sin()).collect();
11085        let (mut am, mut tm) = (vec![0f32; b * rows], vec![0f32; b * rows]);
11086        q4matmat(&split, &xs, b, rows, cols, &mut am, None);
11087        q4t_matmat(&tiled, &xs, b, rows, cols, &mut tm, None);
11088        assert_eq!(am, tm, "q4t matmat must match q4 bit-for-bit");
11089    }
11090
11091    /// q4 SDOT outlier correction: a single huge activation channel
11092    /// (>8·rms → outlier, zeroed in xq) must still contribute its EXACT
11093    /// term. On-grid bulk (±1/0 → xq dequantizes exactly) isolates the
11094    /// correction from A8W8 noise. cols must exceed 64: at n=64 the
11095    /// 8·rms threshold equals sqrt(v²+rest) ≥ v, so a single outlier
11096    /// can never qualify (8² = n).
11097    #[test]
11098    fn q4matvec_sdot_outlier_exact() {
11099        let (rows, cols) = (4, 128);
11100        let groups = rows * cols / GROUP_SIZE;
11101        let mut bytes = Vec::with_capacity(groups * 16 + groups * 2);
11102        for i in 0..groups * 16 {
11103            bytes.push(((i * 11 + 5) % 256) as u8);
11104        }
11105        for g in 0..groups {
11106            let s = 0.02 + 0.002 * g as f32;
11107            bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
11108        }
11109        let mut x: Vec<f32> = (0..cols)
11110            .map(|i| match i % 3 {
11111                0 => 1.0,
11112                1 => -1.0,
11113                _ => 0.0,
11114            })
11115            .collect();
11116        x[17] = 300.0; // ≫ 8·rms → outlier channel
11117
11118        let mut reference = vec![0.0f32; rows * cols];
11119        cortiq_core::quant::dequant_q4_block(&bytes, &mut reference);
11120        let mut expect = vec![0.0f32; rows];
11121        for r in 0..rows {
11122            expect[r] = reference[r * cols..(r + 1) * cols]
11123                .iter()
11124                .zip(&x)
11125                .map(|(w, xv)| w * xv)
11126                .sum();
11127        }
11128        let mut got = vec![0.0f32; rows];
11129        q4matvec(&bytes, &x, rows, cols, &mut got, None);
11130        let scale = expect.iter().fold(0f32, |m, v| m.max(v.abs())).max(1.0);
11131        for r in 0..rows {
11132            assert!(
11133                (got[r] - expect[r]).abs() < 2e-3 * scale,
11134                "row {r}: {} vs {} (outlier term must be exact)",
11135                got[r],
11136                expect[r]
11137            );
11138        }
11139    }
11140
11141    /// The fused q1t matvec must equal the reference (dequant_q1t → dot),
11142    /// including the ternary zero level and the binary-searched outlier
11143    /// overlay. Guards the mmap kernel that makes a 12B q1t runnable.
11144    #[test]
11145    fn q1t_matvec_matches_reference() {
11146        use cortiq_core::quant::{dequant_q1t, f32_to_f16};
11147        let (rows, cols) = (3usize, 64usize); // gpr = 2
11148        let gpr = cols / GROUP_SIZE;
11149        let scales = [0.5f32, 0.3, 0.7, 0.2, 0.6, 0.15];
11150        // Overlay (must be sorted by flat index): a few spikes across rows.
11151        let outliers: [(u32, f32); 3] = [(5, 9.0), (70, -4.5), (150, 3.25)];
11152        let is_out = |flat: usize| outliers.iter().any(|&(i, _)| i as usize == flat);
11153        let mut bytes = Vec::new();
11154        for r in 0..rows {
11155            for g in 0..gpr {
11156                bytes.extend_from_slice(&f32_to_f16(scales[r * gpr + g]).to_le_bytes());
11157                let mut c = [0u8; 7];
11158                for k in 0..GROUP_SIZE {
11159                    // Encoder invariant: code 0 at outlier positions.
11160                    let code = if is_out(r * cols + g * GROUP_SIZE + k) {
11161                        0
11162                    } else {
11163                        ((k + r * 3 + g) % 3) as u8 // 0,1,2
11164                    };
11165                    cortiq_core::quant::q1t_pack(&mut c, k, code);
11166                }
11167                bytes.extend_from_slice(&c);
11168            }
11169        }
11170        // Per-row overlay: [u32 row_ptr[rows+1]] then [(u16 col, f16 val)] by
11171        // row (outliers are sorted by flat index → already grouped by row).
11172        let mut row_ptr = vec![0u32; rows + 1];
11173        for &(idx, _) in &outliers {
11174            row_ptr[idx as usize / cols + 1] += 1;
11175        }
11176        for r in 0..rows {
11177            row_ptr[r + 1] += row_ptr[r];
11178        }
11179        for &p in &row_ptr {
11180            bytes.extend_from_slice(&p.to_le_bytes());
11181        }
11182        for &(idx, v) in &outliers {
11183            bytes.extend_from_slice(&((idx as usize % cols) as u16).to_le_bytes());
11184            bytes.extend_from_slice(&f32_to_f16(v).to_le_bytes());
11185        }
11186
11187        let mut refw = vec![0f32; rows * cols];
11188        dequant_q1t(&bytes, rows, cols, &mut refw);
11189        // On-grid activations (±1, amax 1) so the int8 SDOT path reconstructs
11190        // x exactly and matches the f32 reference (same trick as the q1 test).
11191        let x: Vec<f32> = (0..cols)
11192            .map(|j| if j % 3 == 0 { 1.0 } else { -1.0 })
11193            .collect();
11194        let mut expect = vec![0f32; rows];
11195        for r in 0..rows {
11196            let mut a = 0.0f32;
11197            for j in 0..cols {
11198                a += refw[r * cols + j] * x[j];
11199            }
11200            expect[r] = a;
11201        }
11202        let tol = |e: f32| 1e-3 * e.abs().max(1e-3);
11203        let mut got = vec![0f32; rows];
11204        q1t_matvec(&bytes, &x, rows, cols, &mut got, None);
11205        for r in 0..rows {
11206            assert!(
11207                (got[r] - expect[r]).abs() < tol(expect[r]),
11208                "row {r}: {} vs {}",
11209                got[r],
11210                expect[r]
11211            );
11212        }
11213        // matmat (b=2, f32 decode path) must agree too.
11214        let x2: Vec<f32> = x.iter().chain(x.iter()).copied().collect();
11215        let mut gm = vec![0f32; 2 * rows];
11216        q1t_matmat(&bytes, &x2, 2, rows, cols, &mut gm, None);
11217        for r in 0..rows {
11218            assert!((gm[r] - expect[r]).abs() < tol(expect[r]));
11219            assert!((gm[rows + r] - expect[r]).abs() < tol(expect[r]));
11220        }
11221        // Fused pair (q1t_matvec2) must equal two single matvecs
11222        // bit-for-bit: same unpack, same group order, same f32
11223        // accumulation per stream. Distinct x2 exercises both lanes.
11224        let xb: Vec<f32> = (0..cols)
11225            .map(|j| if j % 5 == 0 { -1.0 } else { 1.0 })
11226            .collect();
11227        let (mut s1, mut s2) = (vec![0f32; rows], vec![0f32; rows]);
11228        q1t_matvec(&bytes, &x, rows, cols, &mut s1, None);
11229        q1t_matvec(&bytes, &xb, rows, cols, &mut s2, None);
11230        let (mut p1, mut p2) = (vec![0f32; rows], vec![0f32; rows]);
11231        q1t_matvec2(&bytes, &x, &xb, rows, cols, &mut p1, &mut p2, None);
11232        assert_eq!(p1, s1, "q1t pair lane 1 ≠ single matvec");
11233        assert_eq!(p2, s2, "q1t pair lane 2 ≠ single matvec");
11234    }
11235
11236    /// Pair == 2×matvec with an ODD group count (the kernel's tail
11237    /// group) and no overlay section.
11238    #[test]
11239    fn q1t_matvec2_odd_gpr_matches_singles() {
11240        use cortiq_core::quant::{Q1T_TILE, f32_to_f16, q1t_pack};
11241        let (rows, cols) = (5usize, 96usize); // gpr = 3 → paired + tail
11242        let gpr = cols / GROUP_SIZE;
11243        let mut bytes = Vec::with_capacity(rows * gpr * Q1T_TILE);
11244        for r in 0..rows {
11245            for g in 0..gpr {
11246                bytes.extend_from_slice(&f32_to_f16(0.1 + 0.05 * (r + g) as f32).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 * 5 + g * 3) % 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        q1t_matvec(&bytes, &x1, rows, cols, &mut s1, None);
11258        q1t_matvec(&bytes, &x2, rows, cols, &mut s2, None);
11259        let (mut p1, mut p2) = (vec![0f32; rows], vec![0f32; rows]);
11260        q1t_matvec2(&bytes, &x1, &x2, rows, cols, &mut p1, &mut p2, None);
11261        assert_eq!(p1, s1, "odd-gpr pair lane 1 ≠ single");
11262        assert_eq!(p2, s2, "odd-gpr pair lane 2 ≠ single");
11263    }
11264
11265    // Speed A/B: fused pair (one unpack, two streams) vs two single
11266    // matvecs. Single-threaded, FFN-sized, min-of paired in-process.
11267    //   cargo test -p cortiq-engine --release q1t_matvec2_speed -- --ignored --nocapture
11268    #[test]
11269    #[ignore]
11270    fn q1t_matvec2_speed() {
11271        use cortiq_core::quant::{Q1T_TILE, f32_to_f16, q1t_pack};
11272        use std::time::Instant;
11273        let (rows, cols) = (8192usize, 4096usize);
11274        let gpr = cols / GROUP_SIZE;
11275        let mut bytes = Vec::with_capacity(rows * gpr * Q1T_TILE);
11276        for r in 0..rows {
11277            for g in 0..gpr {
11278                let s = 0.1 + ((r + g) % 7) as f32 * 0.01;
11279                bytes.extend_from_slice(&f32_to_f16(s).to_le_bytes());
11280                let mut c = [0u8; 7];
11281                for k in 0..GROUP_SIZE {
11282                    q1t_pack(&mut c, k, ((k * 7 + r + g) % 3) as u8);
11283                }
11284                bytes.extend_from_slice(&c);
11285            }
11286        }
11287        let x1: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.31).sin()).collect();
11288        let x2: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.17).cos()).collect();
11289        let (mut s1, mut s2) = (vec![0f32; rows], vec![0f32; rows]);
11290        let (mut p1, mut p2) = (vec![0f32; rows], vec![0f32; rows]);
11291        // Warm both paths once.
11292        q1t_matvec(&bytes, &x1, rows, cols, &mut s1, None);
11293        q1t_matvec2(&bytes, &x1, &x2, rows, cols, &mut p1, &mut p2, None);
11294        let (mut t_pair, mut t_two) = (f64::MAX, f64::MAX);
11295        for _ in 0..8 {
11296            let t0 = Instant::now();
11297            q1t_matvec2(&bytes, &x1, &x2, rows, cols, &mut p1, &mut p2, None);
11298            t_pair = t_pair.min(t0.elapsed().as_secs_f64() * 1000.0);
11299            let t1 = Instant::now();
11300            q1t_matvec(&bytes, &x1, rows, cols, &mut s1, None);
11301            q1t_matvec(&bytes, &x2, rows, cols, &mut s2, None);
11302            t_two = t_two.min(t1.elapsed().as_secs_f64() * 1000.0);
11303        }
11304        assert_eq!(p1, s1);
11305        assert_eq!(p2, s2);
11306        println!("q1t pair {rows}x{cols}: fused {t_pair:.2} ms | two singles {t_two:.2} ms");
11307    }
11308
11309    // Speed A/B: the base-3-division decode (what the packing commit left in
11310    // place) vs the fused sign-LUT matvec. Both single-threaded, same bytes.
11311    //   cargo test -p cortiq-engine q1t_matvec_speed -- --ignored --nocapture
11312    #[test]
11313    #[ignore]
11314    fn q1t_matvec_speed() {
11315        use cortiq_core::quant::{Q1T_TILE, f32_to_f16, q1t_code, q1t_pack};
11316        use std::time::Instant;
11317        let (rows, cols) = (8192usize, 4096usize); // FFN-sized
11318        let gpr = cols / GROUP_SIZE;
11319        let mut bytes = Vec::with_capacity(rows * gpr * Q1T_TILE + 16);
11320        for r in 0..rows {
11321            for g in 0..gpr {
11322                let s = 0.1 + ((r + g) % 7) as f32 * 0.01;
11323                bytes.extend_from_slice(&f32_to_f16(s).to_le_bytes());
11324                let mut c = [0u8; 7];
11325                for k in 0..GROUP_SIZE {
11326                    q1t_pack(&mut c, k, ((k * 7 + r + g) % 3) as u8);
11327                }
11328                bytes.extend_from_slice(&c);
11329            }
11330        }
11331        let (n, stride) = (rows * cols, 40usize); // ~2.5% outliers, per-row overlay
11332        let mut row_ptr = vec![0u32; rows + 1];
11333        let mut idx = 0usize;
11334        while idx < n {
11335            row_ptr[idx / cols + 1] += 1;
11336            idx += stride;
11337        }
11338        for r in 0..rows {
11339            row_ptr[r + 1] += row_ptr[r];
11340        }
11341        for &p in &row_ptr {
11342            bytes.extend_from_slice(&p.to_le_bytes());
11343        }
11344        let mut idx = 0usize;
11345        while idx < n {
11346            bytes.extend_from_slice(&((idx % cols) as u16).to_le_bytes());
11347            bytes.extend_from_slice(&f32_to_f16((idx % 13) as f32 * 0.1 - 0.6).to_le_bytes());
11348            idx += stride;
11349        }
11350        // On-grid ±1 so the fast path's int8 SDOT is exact vs the f32 "slow"
11351        // reference (the A/B is a timing check; values must still agree).
11352        let x: Vec<f32> = (0..cols)
11353            .map(|j| if j % 3 == 0 { 1.0 } else { -1.0 })
11354            .collect();
11355        let (rp_off, ent_off, has_ov) = q1t_overlay(&bytes, rows * gpr * Q1T_TILE, rows);
11356
11357        // "before": base-3 division decode into a buffer, then dot.
11358        let slow = |out: &mut [f32]| {
11359            let mut buf = vec![0f32; cols];
11360            for r in 0..rows {
11361                for g in 0..gpr {
11362                    let off = (r * gpr + g) * Q1T_TILE;
11363                    let s = f16_to_f32(u16::from_le_bytes([bytes[off], bytes[off + 1]]));
11364                    let codes = &bytes[off + 2..off + Q1T_TILE];
11365                    for k in 0..GROUP_SIZE {
11366                        buf[g * GROUP_SIZE + k] = match q1t_code(codes, k) {
11367                            1 => s,
11368                            2 => -s,
11369                            _ => 0.0,
11370                        };
11371                    }
11372                }
11373                out[r] = q1t_row_outlier_correction(&bytes, r, rp_off, ent_off, has_ov, &x)
11374                    + (0..cols).map(|j| buf[j] * x[j]).sum::<f32>();
11375            }
11376        };
11377        let iters = 5;
11378        let mut a = vec![0f32; rows];
11379        slow(&mut a); // warm
11380        let t = Instant::now();
11381        for _ in 0..iters {
11382            slow(&mut a);
11383        }
11384        let slow_ms = t.elapsed().as_secs_f64() * 1e3 / iters as f64;
11385
11386        let mut b = vec![0f32; rows];
11387        q1t_matvec(&bytes, &x, rows, cols, &mut b, None); // warm
11388        let t = Instant::now();
11389        for _ in 0..iters {
11390            q1t_matvec(&bytes, &x, rows, cols, &mut b, None);
11391        }
11392        let fast_ms = t.elapsed().as_secs_f64() * 1e3 / iters as f64;
11393
11394        for r in 0..rows {
11395            assert!((a[r] - b[r]).abs() < 1e-2, "mismatch row {r}");
11396        }
11397        println!(
11398            "q1t matvec {rows}x{cols} (1 thread): div-decode {slow_ms:.2} ms  fused-LUT {fast_ms:.2} ms  => {:.2}x",
11399            slow_ms / fast_ms
11400        );
11401    }
11402}
11403
11404#[cfg(test)]
11405mod gemm_bench {
11406    /// `cargo test -p cortiq-engine --release q4tp_matmat_throughput -- --ignored --nocapture`
11407    /// Times the batched q4tp GEMM at the shapes the image DiT runs
11408    /// (b=296 tokens, 2304 -> 9216), on synthetic bytes: no model, no
11409    /// mmap, no thermal drift over minutes — a kernel change shows up
11410    /// here in seconds where a full render hides it in noise.
11411    ///
11412    /// On macOS add `CMF_ACCEL=0`: this shape is over the 500k-cell mark
11413    /// where the matmat hands off to Accelerate's dequant sgemm, and
11414    /// without the opt-out both rows below measure the AMX, not the
11415    /// kernel under test.
11416    #[test]
11417    #[ignore]
11418    fn q4tp_matmat_throughput() {
11419        // 296 is a prompt-encode batch; the image DiT runs 2085 at
11420        // 512x512, where the activation panel stops fitting L2 and the
11421        // loop's shape starts to matter more than its instructions.
11422        let b: usize = std::env::var("CMF_BENCH_B")
11423            .ok()
11424            .and_then(|v| v.parse().ok())
11425            .unwrap_or(296);
11426        let (rows, cols) = (9216usize, 2304usize);
11427        let (_, _, _) = (rows, cols, b);
11428        let total =
11429            cortiq_core::quant::expected_nbytes(cortiq_core::TensorDtype::Q4TiledP, &[rows, cols])
11430                .unwrap();
11431        // Random nibbles are fine, but the row params are f16 (lo, step)
11432        // of a geometric ladder: garbage there gives exp2 of a huge
11433        // exponent, the scales come back inf, and the whole bench times
11434        // NaN arithmetic instead of the kernel.
11435        let (params_off, codes_off, _) = cortiq_core::quant::q4tp_sections(rows, cols);
11436        let mut bytes: Vec<u8> = (0..total).map(|i| (i * 37 % 251) as u8).collect();
11437        let lo = cortiq_core::quant::f32_to_f16(-4.0);
11438        let step = cortiq_core::quant::f32_to_f16(0.1);
11439        for r in 0..rows {
11440            let o = params_off + r * 4;
11441            bytes[o..o + 2].copy_from_slice(&lo.to_le_bytes());
11442            bytes[o + 2..o + 4].copy_from_slice(&step.to_le_bytes());
11443        }
11444        let _ = codes_off;
11445        let xs: Vec<f32> = (0..b * cols)
11446            .map(|i| ((i % 97) as f32 - 48.0) / 48.0)
11447            .collect();
11448        let mut out = vec![0f32; b * rows];
11449        let pool = crate::pool::Pool::from_env();
11450        // A shared 48-core stand drifts ±25% run to run, which is wider
11451        // than any kernel change worth making. So: alternate the two
11452        // kernels inside one process and keep the BEST time for
11453        // each. Interleaving makes both see the same interference, and a
11454        // minimum is the one statistic another tenant cannot inflate.
11455        super::q4tp_matmat(&bytes, &xs, b, rows, cols, &mut out, pool.as_deref());
11456        let reps: usize = std::env::var("CMF_BENCH_REPS")
11457            .ok()
11458            .and_then(|v| v.parse().ok())
11459            .unwrap_or(10);
11460        let mut best = [f64::MAX; 2];
11461        let mut sums = [0f32; 2];
11462        for _ in 0..reps {
11463            for (k, w) in [(0usize, 1u8), (1usize, 2u8)] {
11464                super::Q4TP_ALT.store(w, std::sync::atomic::Ordering::Relaxed);
11465                let t = std::time::Instant::now();
11466                super::q4tp_matmat(&bytes, &xs, b, rows, cols, &mut out, pool.as_deref());
11467                best[k] = best[k].min(t.elapsed().as_secs_f64());
11468                sums[k] = out.iter().take(64).sum::<f32>();
11469            }
11470        }
11471        let flops = 2.0 * b as f64 * rows as f64 * cols as f64;
11472        for (k, name) in ["previous", "tuned   "].iter().enumerate() {
11473            println!(
11474                "q4tp matmat {rows}x{cols} b={b} {name}: {:.1} ms  {:.1} GFLOP/s  (checksum {:.3})",
11475                best[k] * 1e3,
11476                flops / best[k] / 1e9,
11477                sums[k]
11478            );
11479        }
11480        assert!(
11481            (sums[0] - sums[1]).abs() < 1e-2,
11482            "the tuned kernel changed the result: {} vs {}",
11483            sums[0],
11484            sums[1]
11485        );
11486    }
11487
11488    /// The blocked kernel must agree with the per-column path exactly —
11489    /// same weights, same activation split, only a different instruction
11490    /// mix. Shapes are chosen to hit the awkward cases: a column count
11491    /// that leaves an odd group (the 512-bit kernel does two at a time),
11492    /// and a batch that does not divide by four.
11493    #[test]
11494    fn q4tp_matmat_blocked_matches_scalar() {
11495        use std::sync::atomic::Ordering::Relaxed;
11496        // The last shape carries the image DiT's column count — 2304, so
11497        // 72 groups of accumulation, which is where a reordered sum can
11498        // actually drift — and runs through the thread pool, since the
11499        // blocked path splits rows across workers. Its row count stays
11500        // under 500k cells on purpose: above that, macOS diverts the whole
11501        // matmat to the Accelerate/AMX dequant sgemm and neither kernel
11502        // here would run.
11503        for &(rows, cols, b) in &[
11504            (64usize, 128usize, 7usize),
11505            (33, 96, 4),
11506            (16, 256, 9),
11507            (192, 2304, 37),
11508        ] {
11509            let total = cortiq_core::quant::expected_nbytes(
11510                cortiq_core::TensorDtype::Q4TiledP,
11511                &[rows, cols],
11512            )
11513            .unwrap();
11514            let (params_off, _, _) = cortiq_core::quant::q4tp_sections(rows, cols);
11515            let mut bytes: Vec<u8> = (0..total).map(|i| (i * 61 % 251) as u8).collect();
11516            let lo = cortiq_core::quant::f32_to_f16(-4.0);
11517            let step = cortiq_core::quant::f32_to_f16(0.1);
11518            for r in 0..rows {
11519                let o = params_off + r * 4;
11520                bytes[o..o + 2].copy_from_slice(&lo.to_le_bytes());
11521                bytes[o + 2..o + 4].copy_from_slice(&step.to_le_bytes());
11522            }
11523            let xs: Vec<f32> = (0..b * cols)
11524                .map(|i| ((i % 89) as f32 - 44.0) / 44.0)
11525                .collect();
11526            let mut got = vec![0f32; b * rows];
11527            let mut want = vec![0f32; b * rows];
11528            let gpr = cols / 32;
11529            let view = super::Q4tpView::new(&bytes, rows, cols);
11530            let pool = crate::pool::Pool::from_env();
11531            super::Q4TP_ALT.store(2, Relaxed);
11532            super::q4tp_matmat(&bytes, &xs, b, rows, cols, &mut got, pool.as_deref());
11533            super::Q4TP_ALT.store(1, Relaxed);
11534            super::q4tp_matmat(&bytes, &xs, b, rows, cols, &mut want, pool.as_deref());
11535            super::Q4TP_ALT.store(0, Relaxed);
11536            // Measured against the output's scale, not cell by cell: a
11537            // dot product of 2304 terms lands near zero wherever the row
11538            // and the activation nearly cancel, and there a per-cell
11539            // ratio reports 1e-3 for an absolute error of 5e-6 — f32's
11540            // own rounding, reordered. What must stay small is the error
11541            // relative to what the layer actually outputs.
11542            let scale = want.iter().fold(0f32, |m, v| m.max(v.abs())).max(1e-6);
11543            let (mut worst, mut at) = (0f32, 0usize);
11544            for (i, (g, w)) in got.iter().zip(&want).enumerate() {
11545                if (g - w).abs() > worst {
11546                    worst = (g - w).abs();
11547                    at = i;
11548                }
11549            }
11550            assert!(
11551                worst <= 1e-4 * scale,
11552                "{rows}x{cols} b={b}: blocked and scalar disagree by {worst:.3e} \
11553                 (scale {scale:.3e}) at cell {at}: {} vs {}",
11554                got[at],
11555                want[at]
11556            );
11557
11558            // "Same speed, no quality loss" is a claim about which answer
11559            // is RIGHT, not about which two agree. Both paths sum the same
11560            // 2304 products in different orders, so f64 decides: the
11561            // blocked kernel keeps sixteen partial sums and folds them at
11562            // the end, which is a shallower addition tree than the
11563            // per-column path's running scalar, and it must not be worse.
11564            let (mut e_blocked, mut e_scalar) = (0f64, 0f64);
11565            for bi in 0..b {
11566                let act = super::split_act(&xs[bi * cols..(bi + 1) * cols]);
11567                for r in 0..rows {
11568                    let mut sc = vec![0f32; gpr];
11569                    view.scales_into(r, gpr, &mut sc);
11570                    let mut exact = 0f64;
11571                    for j in 0..cols {
11572                        let (w, sq) = super::q4tp_outlier(view.nib, r, gpr, j, &sc);
11573                        exact += w as f64 * sq as f64 * act.xq[j] as f64;
11574                    }
11575                    exact *= act.sx as f64;
11576                    for &(j, xv) in &act.outliers {
11577                        let (w, sq) = super::q4tp_outlier(view.nib, r, gpr, j, &sc);
11578                        exact += w as f64 * sq as f64 * xv as f64;
11579                    }
11580                    let i = bi * rows + r;
11581                    e_blocked = e_blocked.max((got[i] as f64 - exact).abs());
11582                    e_scalar = e_scalar.max((want[i] as f64 - exact).abs());
11583                }
11584            }
11585            println!(
11586                "{rows}x{cols} b={b}: worst error vs f64 — blocked {e_blocked:.3e}, \
11587                 per-column {e_scalar:.3e}"
11588            );
11589            // An absolute bar, not a race between the two: at these
11590            // magnitudes both sit in f32's last bits, and on a small shape
11591            // whichever one happens to round the unluckiest cell "wins" by
11592            // a factor the next seed reverses.
11593            assert!(
11594                e_blocked <= 1e-5 * scale as f64 && e_scalar <= 1e-5 * scale as f64,
11595                "{rows}x{cols} b={b}: error against f64 too large — blocked \
11596                 {e_blocked:.3e}, per-column {e_scalar:.3e}, scale {scale:.3e}"
11597            );
11598        }
11599    }
11600
11601    /// The q4t twin of the throughput bench, same shape and rules, so the
11602    /// two quantisations' batch kernels can be read against each other.
11603    /// `cargo test -p cortiq-engine --release q4t_matmat_throughput -- --ignored --nocapture`
11604    #[test]
11605    #[ignore]
11606    fn q4t_matmat_throughput() {
11607        let (rows, cols, b) = (9216usize, 2304usize, 296usize);
11608        let total =
11609            cortiq_core::quant::expected_nbytes(cortiq_core::TensorDtype::Q4Tiled, &[rows, cols])
11610                .unwrap();
11611        // q4t carries a per-group f16 scale in the tile's first two bytes;
11612        // random bytes there decode to inf and the bench would time NaNs.
11613        let mut bytes: Vec<u8> = (0..total).map(|i| (i * 37 % 251) as u8).collect();
11614        let sc = cortiq_core::quant::f32_to_f16(0.02);
11615        for t in bytes.chunks_mut(super::Q4_TILE) {
11616            t[..2].copy_from_slice(&sc.to_le_bytes());
11617        }
11618        let xs: Vec<f32> = (0..b * cols)
11619            .map(|i| ((i % 97) as f32 - 48.0) / 48.0)
11620            .collect();
11621        let mut out = vec![0f32; b * rows];
11622        let pool = crate::pool::Pool::from_env();
11623        super::q4t_matmat(&bytes, &xs, b, rows, cols, &mut out, pool.as_deref());
11624        let reps: usize = std::env::var("CMF_BENCH_REPS")
11625            .ok()
11626            .and_then(|v| v.parse().ok())
11627            .unwrap_or(10);
11628        let mut best = f64::MAX;
11629        for _ in 0..reps {
11630            let t = std::time::Instant::now();
11631            super::q4t_matmat(&bytes, &xs, b, rows, cols, &mut out, pool.as_deref());
11632            best = best.min(t.elapsed().as_secs_f64());
11633        }
11634        let flops = 2.0 * b as f64 * rows as f64 * cols as f64;
11635        println!(
11636            "q4t matmat {rows}x{cols} b={b}: {:.1} ms  {:.1} GFLOP/s  (checksum {:.3})",
11637            best * 1e3,
11638            flops / best / 1e9,
11639            out.iter().take(64).sum::<f32>()
11640        );
11641    }
11642}