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 mapped weight in ANY codec the fused device
526    /// paths can run — four-bit tiled or either int8 layout.
527    ///
528    /// The fused DiT chains asked for `mapped_q4tp` by name, so an eight-bit
529    /// container never reached them and rendered through per-op GEMMs even
530    /// after those kernels learned its codec. The gate is what the codec has
531    /// a device GEMM for, not which codec it is.
532    pub fn mapped_device_gemm(&self) -> Option<(&Arc<CmfModel>, usize)> {
533        match self {
534            Self::Mapped {
535                model,
536                idx,
537                dtype: TensorDtype::Q4TiledP | TensorDtype::Q8Row | TensorDtype::Q8_2f,
538                ..
539            } => Some((model, *idx)),
540            _ => None,
541        }
542    }
543
544    /// (model, tensor idx) for a q2tp mapped weight — the 2-bit twin of
545    /// `mapped_q4tp`, used by the mixed MoE profile.
546    pub fn mapped_q2tp(&self) -> Option<(&Arc<CmfModel>, usize)> {
547        match self {
548            Self::Mapped {
549                model,
550                idx,
551                dtype: TensorDtype::Q2TiledP,
552                ..
553            } => Some((model, *idx)),
554            _ => None,
555        }
556    }
557
558    pub fn cols(&self) -> usize {
559        match self {
560            Self::F32 { cols, .. } | Self::Mapped { cols, .. } => *cols,
561        }
562    }
563
564    /// (model, tensor idx) for a q1 mapped weight — the wgpu token graph
565    /// keys its resident VRAM cache by idx. None for any other dtype/kind.
566    pub fn mapped_q1(&self) -> Option<(&std::sync::Arc<CmfModel>, usize)> {
567        match self {
568            Self::Mapped {
569                model,
570                idx,
571                dtype: TensorDtype::Q1,
572                ..
573            } => Some((model, *idx)),
574            _ => None,
575        }
576    }
577
578    /// (model, idx, kind, row_scale) for a graph-capable mapped weight.
579    /// kind: 0=q8_row (per-row scales), 1=q1, 2=q4_block, 3=q1t
580    /// (tile-embedded, no rs), 5=q4_tiled, 6=q4tp, 7=q8_2f (both scale
581    /// planes live inside the tensor). None only for `vbit`.
582    ///
583    /// The old comment here claimed q4_block was unhandled while the arm
584    /// right below mapped it, and it named q8_2f as unhandled after that
585    /// stopped being true — a stale comment on this function is how a
586    /// model silently loses the graph, so it is worth keeping honest.
587    pub fn graph_weight(&self) -> Option<(&std::sync::Arc<CmfModel>, usize, u8, &[f32])> {
588        match self {
589            Self::Mapped {
590                model,
591                idx,
592                dtype: TensorDtype::Q8Row,
593                row_scale,
594                ..
595            } => Some((model, *idx, 0, row_scale.as_slice())),
596            Self::Mapped {
597                model,
598                idx,
599                dtype: TensorDtype::Q1,
600                ..
601            } => Some((model, *idx, 1, &[])),
602            // Q4Tiled is kind 5, NOT 2: both carried 2 historically, and
603            // the wgpu token graph fed 18B interleaved tiles to the
604            // split-layout q4b kernel — garbage output on q4t models
605            // (caught by an end-to-end answer check on real Vulkan).
606            Self::Mapped {
607                model,
608                idx,
609                dtype: TensorDtype::Q4Tiled,
610                ..
611            } => Some((model, *idx, 5, &[])),
612            // Kind 6, not 5: q4tp's nibble stride and scale planes differ,
613            // and feeding them to the q4t kernel is exactly the mistake that
614            // produced garbage when Q4Tiled shared kind 2 with Q4Block.
615            Self::Mapped {
616                model,
617                idx,
618                dtype: TensorDtype::Q4TiledP,
619                ..
620            } => Some((model, *idx, 6, &[])),
621            Self::Mapped {
622                model,
623                idx,
624                dtype: TensorDtype::Q4Block,
625                ..
626            } => Some((model, *idx, 2, &[])),
627            // q8_2f carries BOTH scale planes after the int8 body (rows
628            // f16, then cols f16), so the graph takes the whole tensor
629            // and the kernel reads them where they lie — no host-side
630            // prescale, which is what the per-op path does instead.
631            Self::Mapped {
632                model,
633                idx,
634                dtype: TensorDtype::Q8_2f,
635                ..
636            } => Some((model, *idx, 7, &[])),
637            Self::Mapped {
638                model,
639                idx,
640                dtype: TensorDtype::Q1T,
641                ..
642            } => Some((model, *idx, 3, &[])),
643            // Kind 9: the 2-bit plane on the q4tp ladder (dense FFN gate/up
644            // of the q2tp profile). Its own kernel — 8 bytes a group where
645            // q4tp has 16, and rung 0 is the exact zero.
646            Self::Mapped {
647                model,
648                idx,
649                dtype: TensorDtype::Q2TiledP,
650                ..
651            } => Some((model, *idx, 9, &[])),
652            _ => None,
653        }
654    }
655
656    /// Dense f32 view — only for owned tensors. Masked/sparse execution
657    /// paths require it; quantized weights don't support masks yet.
658    pub fn as_f32(&self) -> Option<&[f32]> {
659        match self {
660            Self::F32 { data, .. } => Some(data),
661            Self::Mapped { .. } => None,
662        }
663    }
664
665    fn quant_bytes(&self) -> &[u8] {
666        match self {
667            Self::Mapped { model, idx, .. } => model.entry_bytes(&model.tensors[*idx]),
668            Self::F32 { .. } => unreachable!("quant_bytes on F32"),
669        }
670    }
671
672    /// Dequantize one row into `dst` (embedding lookup).
673    pub fn row_f32(&self, r: usize, dst: &mut [f32]) {
674        let cols = self.cols();
675        debug_assert_eq!(dst.len(), cols);
676        match self {
677            Self::F32 { data, .. } => dst.copy_from_slice(&data[r * cols..(r + 1) * cols]),
678            Self::Mapped {
679                dtype,
680                row_scale,
681                col_field,
682                vbit_offsets,
683                ..
684            } => {
685                if *dtype == TensorDtype::Q4Tiled {
686                    let bytes = self.quant_bytes();
687                    let gpr = cols / GROUP_SIZE;
688                    for gi in 0..gpr {
689                        let tile = &bytes[(r * gpr + gi) * Q4_TILE..(r * gpr + gi + 1) * Q4_TILE];
690                        let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
691                        for (k, &b) in tile[2..].iter().enumerate() {
692                            dst[gi * GROUP_SIZE + k * 2] = ((b & 0x0F) as f32 - 8.0) * s;
693                            dst[gi * GROUP_SIZE + k * 2 + 1] = (((b >> 4) & 0x0F) as f32 - 8.0) * s;
694                        }
695                    }
696                    return;
697                }
698                if *dtype == TensorDtype::Q4TiledP {
699                    let bytes = self.quant_bytes();
700                    let gpr = cols / GROUP_SIZE;
701                    let v = Q4tpView::new(bytes, self.rows(), cols);
702                    let mut sc = vec![0f32; gpr];
703                    v.scales_into(r, gpr, &mut sc);
704                    for gi in 0..gpr {
705                        let tile = &v.nib[(r * gpr + gi) * Q4TP_NIB..(r * gpr + gi + 1) * Q4TP_NIB];
706                        let s = sc[gi];
707                        for (k, &b) in tile.iter().enumerate() {
708                            dst[gi * GROUP_SIZE + k * 2] = ((b & 0x0F) as f32 - 8.0) * s;
709                            dst[gi * GROUP_SIZE + k * 2 + 1] = (((b >> 4) & 0x0F) as f32 - 8.0) * s;
710                        }
711                    }
712                    return;
713                }
714                if *dtype == TensorDtype::Q2TiledP {
715                    let bytes = self.quant_bytes();
716                    let gpr = cols / GROUP_SIZE;
717                    let v = Q4tpView::new_q2(bytes, self.rows(), cols);
718                    let mut sc = vec![0f32; gpr];
719                    v.scales_into(r, gpr, &mut sc);
720                    for gi in 0..gpr {
721                        let ch =
722                            &v.nib[(r * gpr + gi) * Q2TP_CHUNK..(r * gpr + gi + 1) * Q2TP_CHUNK];
723                        let s = sc[gi];
724                        for (k, &b) in ch.iter().enumerate() {
725                            for j in 0..4 {
726                                dst[gi * GROUP_SIZE + k * 4 + j] =
727                                    (((b >> (2 * j)) & 3) as f32 - 1.5) * s;
728                            }
729                        }
730                    }
731                    return;
732                }
733                if *dtype == TensorDtype::Q4Block {
734                    let (packed, scales) = q4_split(self.quant_bytes(), self.rows(), cols);
735                    let gpr = cols / GROUP_SIZE;
736                    for gi in 0..gpr {
737                        let g = r * gpr + gi;
738                        let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
739                        for (k, &b) in packed[g * 16..(g + 1) * 16].iter().enumerate() {
740                            dst[gi * GROUP_SIZE + k * 2] = ((b & 0x0F) as f32 - 8.0) * s;
741                            dst[gi * GROUP_SIZE + k * 2 + 1] = (((b >> 4) & 0x0F) as f32 - 8.0) * s;
742                        }
743                    }
744                    return;
745                }
746                if *dtype == TensorDtype::Q1 {
747                    let bytes = self.quant_bytes();
748                    let gpr = cols / GROUP_SIZE;
749                    for gi in 0..gpr {
750                        let tile = &bytes[(r * gpr + gi) * Q1_TILE..(r * gpr + gi + 1) * Q1_TILE];
751                        let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
752                        for (j, &b) in tile[2..].iter().enumerate() {
753                            for k in 0..8 {
754                                dst[gi * GROUP_SIZE + j * 8 + k] =
755                                    (((b >> k) & 1) as f32 * 2.0 - 1.0) * s;
756                            }
757                        }
758                    }
759                    return;
760                }
761                if *dtype == TensorDtype::Q1T {
762                    let bytes = self.quant_bytes();
763                    let gpr = cols / GROUP_SIZE;
764                    let base_len = self.rows() * gpr * cortiq_core::quant::Q1T_TILE;
765                    for gi in 0..gpr {
766                        let off = (r * gpr + gi) * cortiq_core::quant::Q1T_TILE;
767                        let s = cortiq_core::quant::f16_to_f32(u16::from_le_bytes([
768                            bytes[off],
769                            bytes[off + 1],
770                        ]));
771                        let codes = &bytes[off + 2..off + cortiq_core::quant::Q1T_TILE];
772                        for k in 0..GROUP_SIZE {
773                            dst[gi * GROUP_SIZE + k] = match cortiq_core::quant::q1t_code(codes, k)
774                            {
775                                1 => s,
776                                2 => -s,
777                                _ => 0.0,
778                            };
779                        }
780                    }
781                    // Overlay
782                    let rows = self.rows();
783                    let entries = base_len + (rows + 1) * 4;
784                    if entries <= bytes.len() {
785                        let ptrs = &bytes[base_len..base_len + (rows + 1) * 4];
786                        let r0 = u32::from_le_bytes([
787                            ptrs[r * 4],
788                            ptrs[r * 4 + 1],
789                            ptrs[r * 4 + 2],
790                            ptrs[r * 4 + 3],
791                        ]) as usize;
792                        let r1 = u32::from_le_bytes([
793                            ptrs[(r + 1) * 4],
794                            ptrs[(r + 1) * 4 + 1],
795                            ptrs[(r + 1) * 4 + 2],
796                            ptrs[(r + 1) * 4 + 3],
797                        ]) as usize;
798                        let off = entries + r0 * 4;
799                        for i in 0..r1 - r0 {
800                            let item = &bytes[off + i * 4..off + i * 4 + 4];
801                            let c = u16::from_le_bytes([item[0], item[1]]) as usize;
802                            let v = cortiq_core::quant::f16_to_f32(u16::from_le_bytes([
803                                item[2], item[3],
804                            ]));
805                            if c < cols {
806                                dst[c] = v;
807                            }
808                        }
809                    }
810                    return;
811                }
812                if matches!(dtype, TensorDtype::Vbit | TensorDtype::VbitRo) {
813                    let bytes = self.quant_bytes();
814                    let rows = self.rows();
815                    let ng = cols / GROUP_SIZE;
816                    let bits = &bytes[..rows];
817                    let sc_off = rows;
818                    // Precomputed at load — embedding lookup used to scan
819                    // the bit-widths of every preceding row (O(token_id)).
820                    let off = vbit_offsets[r];
821                    let b = bits[r] as usize;
822                    let l = ((1usize << (b - 1)) - 1) as f32;
823                    let data = &bytes[off..];
824                    let (mut acc, mut nbits, mut idx) = (0u64, 0usize, 0usize);
825                    for (i, d) in dst.iter_mut().enumerate() {
826                        while nbits < b {
827                            acc = (acc << 8) | data[idx] as u64;
828                            idx += 1;
829                            nbits += 8;
830                        }
831                        let u = ((acc >> (nbits - b)) & ((1u64 << b) - 1)) as f32;
832                        nbits -= b;
833                        let so = (r * ng + i / GROUP_SIZE) * 2;
834                        let sv = f16_to_f32(u16::from_le_bytes([
835                            bytes[sc_off + so],
836                            bytes[sc_off + so + 1],
837                        ]));
838                        *d = (u - l) * sv;
839                    }
840                    return;
841                }
842                let q = &self.quant_bytes()[r * cols..(r + 1) * cols];
843                let s = row_scale[r];
844                match dtype {
845                    TensorDtype::Q8Row => {
846                        for (d, &b) in dst.iter_mut().zip(q) {
847                            *d = (b as i8) as f32 * s;
848                        }
849                    }
850                    TensorDtype::Q8_2f => {
851                        for (i, (d, &b)) in dst.iter_mut().zip(q).enumerate() {
852                            *d = (b as i8) as f32 * s * col_field[i];
853                        }
854                    }
855                    _ => unreachable!(),
856                }
857            }
858        }
859    }
860
861    /// Can this tensor's columns be read cheaply (for sparse down_proj)?
862    /// True for F32/Q8Row/Q8_2f (per-row scale, direct strided access);
863    /// false for group-packed q4/vbit (column access would unpack whole
864    /// groups — sparse execution falls back to f32 for those).
865    pub fn sparse_col_ok(&self) -> bool {
866        match self {
867            Self::F32 { .. } => true,
868            Self::Mapped { dtype, .. } => {
869                matches!(dtype, TensorDtype::Q8Row | TensorDtype::Q8_2f)
870            }
871        }
872    }
873
874    /// down_proj [hidden, inter]: accumulate `w · col(c)` into `out`
875    /// [hidden] — reads ONLY column `c` (one neuron) from the mmap,
876    /// no full-matrix dequant. `out[k] += w · down[k, c]`.
877    pub fn add_col_scaled(&self, c: usize, w: f32, out: &mut [f32]) {
878        let inter = self.cols();
879        let hidden = self.rows();
880        debug_assert_eq!(out.len(), hidden);
881        match self {
882            Self::F32 { data, .. } => {
883                for (k, o) in out.iter_mut().enumerate() {
884                    *o += w * data[k * inter + c];
885                }
886            }
887            Self::Mapped {
888                dtype,
889                row_scale,
890                col_field,
891                ..
892            } => {
893                let q = self.quant_bytes();
894                let colf = if *dtype == TensorDtype::Q8_2f {
895                    col_field[c]
896                } else {
897                    1.0
898                };
899                let wc = w * colf;
900                for (k, o) in out.iter_mut().enumerate() {
901                    let b = q[k * inter + c] as i8 as f32;
902                    *o += wc * b * row_scale[k];
903                }
904            }
905        }
906    }
907
908    /// Touch the head of row `r` so the DRAM latency of the next
909    /// neuron's weights overlaps the current one's arithmetic.
910    ///
911    /// Scattered rows are what per-token sparsity reads, and a 2 KB
912    /// stride is past what the hardware prefetcher follows: without this
913    /// every row starts with a cold miss that nothing hides. One touch
914    /// per 512 bytes is enough — the rest of the row is a sequential run
915    /// the prefetcher does pick up.
916    #[inline]
917    pub fn prefetch_row(&self, r: usize) {
918        let Self::Mapped { dtype, .. } = self else {
919            return;
920        };
921        if !matches!(dtype, TensorDtype::Q8Row | TensorDtype::Q8_2f) {
922            return;
923        }
924        let cols = self.cols();
925        let q = self.quant_bytes();
926        let (a, b) = (r * cols, (r + 1) * cols);
927        if b > q.len() {
928            return;
929        }
930        let mut j = a;
931        while j < b {
932            unsafe { std::ptr::read_volatile(q.as_ptr().add(j)) };
933            j += 512;
934        }
935    }
936
937    /// `out += w · row(r)` — the transposed twin of `add_col_scaled`.
938    ///
939    /// A neuron's `down` weights are a COLUMN of `[hidden, inter]`, and a
940    /// column is strided: reading one costs a cache line per element, so
941    /// per-neuron dynamic sparsity saves arithmetic and no bytes. Stored
942    /// transposed (`down_proj.t.weight`, `[inter, hidden]`) the same
943    /// weights are a contiguous ROW, and this accumulate reads exactly
944    /// the neurons the token asked for.
945    pub fn add_row_scaled(&self, r: usize, w: f32, out: &mut [f32], scratch: &mut [f32]) {
946        let cols = self.cols();
947        debug_assert_eq!(out.len(), cols);
948        match self {
949            Self::F32 { data, .. } => {
950                let row = &data[r * cols..(r + 1) * cols];
951                for (o, v) in out.iter_mut().zip(row) {
952                    *o += w * v;
953                }
954            }
955            Self::Mapped {
956                dtype,
957                row_scale,
958                col_field,
959                ..
960            } => match dtype {
961                TensorDtype::Q8Row => {
962                    let q = &self.quant_bytes()[r * cols..(r + 1) * cols];
963                    let ws = w * row_scale[r];
964                    let row: &[i8] =
965                        unsafe { std::slice::from_raw_parts(q.as_ptr() as *const i8, q.len()) };
966                    axpy_i8_f32(out, row, ws);
967                }
968                TensorDtype::Q8_2f => {
969                    let q = &self.quant_bytes()[r * cols..(r + 1) * cols];
970                    let ws = w * row_scale[r];
971                    for ((o, b), c) in out.iter_mut().zip(q).zip(col_field) {
972                        *o += ws * c * (*b as i8 as f32);
973                    }
974                }
975                _ => {
976                    self.row_f32(r, scratch);
977                    for (o, v) in out.iter_mut().zip(scratch.iter()) {
978                        *o += w * v;
979                    }
980                }
981            },
982        }
983    }
984
985    /// Dot of row `r` with `x` (gate/up active-neuron path). Reads only
986    /// row `r` from the mmap — no full dequant. q4/vbit dequant the row
987    /// into `scratch` first (rare for active-FFN weights).
988    pub fn row_dot(&self, r: usize, x: &[f32], scratch: &mut [f32]) -> f32 {
989        let cols = self.cols();
990        match self {
991            Self::F32 { data, .. } => {
992                let row = &data[r * cols..(r + 1) * cols];
993                row.iter().zip(x).map(|(w, v)| w * v).sum()
994            }
995            Self::Mapped {
996                dtype,
997                row_scale,
998                col_field,
999                ..
1000            } => match dtype {
1001                TensorDtype::Q8Row => {
1002                    let q = &self.quant_bytes()[r * cols..(r + 1) * cols];
1003                    dot_i8_f32(q, x) * row_scale[r]
1004                }
1005                TensorDtype::Q8_2f => {
1006                    let q = &self.quant_bytes()[r * cols..(r + 1) * cols];
1007                    dot_i8_col_f32(q, x, col_field) * row_scale[r]
1008                }
1009                _ => {
1010                    self.row_f32(r, scratch);
1011                    scratch.iter().zip(x).map(|(w, v)| w * v).sum()
1012                }
1013            },
1014        }
1015    }
1016
1017    /// `out = W · x` (row-major). F32 delegates to the historical
1018    /// bit-exact path; Mapped runs the fused int8 kernel.
1019    pub fn matvec(&self, x: &[f32], out: &mut [f32], pool: Option<&Pool>) {
1020        match self {
1021            // NOTE: `out.len()` DRIVES this arm — it computes that many rows,
1022            // and `x.len()` is the stride. A short `out` is legitimate here,
1023            // which is why the check below lives in the Mapped arm only.
1024            Self::F32 { data, .. } => matvec_rows(pool, data, x, out),
1025            Self::Mapped {
1026                model,
1027                idx,
1028                dtype,
1029                rows,
1030                cols,
1031                row_scale,
1032                col_field,
1033                vbit_offsets,
1034                repack,
1035            } => {
1036                let _ = (model, idx);
1037                // Every kernel below writes `rows` entries through a raw
1038                // pointer, so a short `out` is an out-of-bounds WRITE, not a
1039                // wrong answer: it scribbles on the allocator's metadata and
1040                // the process aborts much later, somewhere innocent
1041                // (`double free or corruption`, `corrupted double-linked
1042                // list`). The debug_assert two of the kernels carried is
1043                // compiled out of the release — exactly the build where it
1044                // matters. Fail here instead, while the caller is still on
1045                // the stack to be named.
1046                assert!(
1047                    out.len() >= *rows && x.len() >= *cols,
1048                    "matvec {rows}x{cols}: out {} (need {rows}), x {} (need {cols})",
1049                    out.len(),
1050                    x.len(),
1051                );
1052                if *dtype == TensorDtype::Q4Block {
1053                    // GPU route (wgpu q4b kernel) for large q4_block matvecs —
1054                    // gives NVIDIA/AMD/Intel q4 models a GPU path. Probe keeps
1055                    // the winner; Metal returns false → the CPU kernel below.
1056                    if *rows * *cols >= 8_388_608 && crate::gpu::enabled_here() {
1057                        let t0 = std::time::Instant::now();
1058                        match crate::gpu::probe_arm(crate::gpu::OpClass::Matvec) {
1059                            crate::gpu::ProbeArm::Gpu => {
1060                                if crate::gpu::q4b_matvec(model, *idx, x, *rows, *cols, out) {
1061                                    crate::gpu::probe_record(
1062                                        crate::gpu::OpClass::Matvec,
1063                                        true,
1064                                        t0.elapsed(),
1065                                    );
1066                                    return;
1067                                }
1068                            }
1069                            crate::gpu::ProbeArm::CpuTimed => {
1070                                q4matvec(self.quant_bytes(), x, *rows, *cols, out, pool);
1071                                crate::gpu::probe_record(
1072                                    crate::gpu::OpClass::Matvec,
1073                                    false,
1074                                    t0.elapsed(),
1075                                );
1076                                return;
1077                            }
1078                            crate::gpu::ProbeArm::Cpu => {}
1079                        }
1080                    }
1081                    q4matvec(self.quant_bytes(), x, *rows, *cols, out, pool);
1082                    return;
1083                }
1084                if *dtype == TensorDtype::Q4Tiled {
1085                    // GPU route for large q4t matvecs — the lm_head class,
1086                    // same shape as the q4tp arm below. The probe keeps the
1087                    // winner; a backend without the kernel refuses and the
1088                    // CPU path stays.
1089                    if *rows * *cols >= 8_388_608 && crate::gpu::enabled_here() {
1090                        let t0 = std::time::Instant::now();
1091                        let cls = crate::gpu::matvec_class(*rows, *cols);
1092                        match crate::gpu::probe_arm(cls) {
1093                            crate::gpu::ProbeArm::Gpu => {
1094                                if crate::gpu::q4t_matvec(model, *idx, x, *rows, *cols, out) {
1095                                    crate::gpu::probe_record(cls, true, t0.elapsed());
1096                                    return;
1097                                }
1098                            }
1099                            crate::gpu::ProbeArm::CpuTimed => {
1100                                q4t_matvec(self.quant_bytes(), x, *rows, *cols, out, pool);
1101                                crate::gpu::probe_record(cls, false, t0.elapsed());
1102                                return;
1103                            }
1104                            crate::gpu::ProbeArm::Cpu => {}
1105                        }
1106                    }
1107                    q4t_matvec(self.quant_bytes(), x, *rows, *cols, out, pool);
1108                    return;
1109                }
1110                if *dtype == TensorDtype::Q4TiledP {
1111                    // GPU route for large q4tp matvecs — the lm_head class.
1112                    // On a q4tp checkpoint the head is the biggest single
1113                    // host matvec left in the decode step, and the batched
1114                    // kernel at b=1 already exists on both backends. Probe
1115                    // keeps the winner, same as q4_block above.
1116                    if *rows * *cols >= 8_388_608 && crate::gpu::enabled_here() {
1117                        let t0 = std::time::Instant::now();
1118                        let cls = crate::gpu::matvec_class(*rows, *cols);
1119                        match crate::gpu::probe_arm(cls) {
1120                            crate::gpu::ProbeArm::Gpu => {
1121                                if crate::gpu::q4tp_matvec(model, *idx, x, *rows, *cols, out) {
1122                                    crate::gpu::probe_record(cls, true, t0.elapsed());
1123                                    return;
1124                                }
1125                            }
1126                            crate::gpu::ProbeArm::CpuTimed => {
1127                                q4tp_matvec(self.quant_bytes(), x, *rows, *cols, out, pool);
1128                                crate::gpu::probe_record(cls, false, t0.elapsed());
1129                                return;
1130                            }
1131                            crate::gpu::ProbeArm::Cpu => {}
1132                        }
1133                    }
1134                    q4tp_matvec(self.quant_bytes(), x, *rows, *cols, out, pool);
1135                    return;
1136                }
1137                if *dtype == TensorDtype::Q2TiledP {
1138                    q2tp_matvec(self.quant_bytes(), x, *rows, *cols, out, pool);
1139                    return;
1140                }
1141                if *dtype == TensorDtype::Q1 {
1142                    // GPU route for large q1 matvecs (out_proj / lm_head
1143                    // class): the CPU q1 kernel is load-port-bound at
1144                    // ~4 GB/s/core, the GPU one is bandwidth-bound — the
1145                    // probe measures both arms and keeps the winner.
1146                    if *rows * *cols >= 8_388_608 && crate::gpu::enabled_here() {
1147                        let t0 = std::time::Instant::now();
1148                        let arm = if crate::gpu::q1_force() {
1149                            crate::gpu::ProbeArm::Gpu
1150                        } else {
1151                            crate::gpu::probe_arm(crate::gpu::OpClass::Matvec)
1152                        };
1153                        match arm {
1154                            crate::gpu::ProbeArm::Gpu => {
1155                                if crate::gpu::q1_matvec(model, *idx, x, *rows, *cols, out) {
1156                                    crate::gpu::probe_record(
1157                                        crate::gpu::OpClass::Matvec,
1158                                        true,
1159                                        t0.elapsed(),
1160                                    );
1161                                    return;
1162                                }
1163                            }
1164                            crate::gpu::ProbeArm::CpuTimed => {
1165                                q1_matvec(self.quant_bytes(), x, *rows, *cols, out, pool);
1166                                crate::gpu::probe_record(
1167                                    crate::gpu::OpClass::Matvec,
1168                                    false,
1169                                    t0.elapsed(),
1170                                );
1171                                return;
1172                            }
1173                            crate::gpu::ProbeArm::Cpu => {}
1174                        }
1175                    }
1176                    q1_matvec(self.quant_bytes(), x, *rows, *cols, out, pool);
1177                    return;
1178                }
1179                if *dtype == TensorDtype::Q1T {
1180                    // GPU route for large q1t matvecs: the ternary BASE dot runs
1181                    // on the GPU (load-port-bound on CPU, like q1), then the
1182                    // sparse overlay is added on the CPU. Probe keeps the winner.
1183                    if *rows * *cols >= 8_388_608 && crate::gpu::enabled_here() {
1184                        let t0 = std::time::Instant::now();
1185                        match crate::gpu::probe_arm(crate::gpu::OpClass::Matvec) {
1186                            crate::gpu::ProbeArm::Gpu => {
1187                                if crate::gpu::q1t_matvec(model, *idx, x, *rows, *cols, out) {
1188                                    q1t_add_overlay(self.quant_bytes(), x, *rows, *cols, out, pool);
1189                                    crate::gpu::probe_record(
1190                                        crate::gpu::OpClass::Matvec,
1191                                        true,
1192                                        t0.elapsed(),
1193                                    );
1194                                    return;
1195                                }
1196                            }
1197                            crate::gpu::ProbeArm::CpuTimed => {
1198                                q1t_matvec(self.quant_bytes(), x, *rows, *cols, out, pool);
1199                                crate::gpu::probe_record(
1200                                    crate::gpu::OpClass::Matvec,
1201                                    false,
1202                                    t0.elapsed(),
1203                                );
1204                                return;
1205                            }
1206                            crate::gpu::ProbeArm::Cpu => {}
1207                        }
1208                    }
1209                    q1t_matvec(self.quant_bytes(), x, *rows, *cols, out, pool);
1210                    return;
1211                }
1212                if matches!(dtype, TensorDtype::Vbit | TensorDtype::VbitRo) {
1213                    vbitmatvec(self.quant_bytes(), vbit_offsets, x, *rows, *cols, out, pool);
1214                    return;
1215                }
1216                let xs = prescale(x, col_field, *dtype);
1217                // D5: large q8 matrices (lm_head-class) — hybrid
1218                // CPU∥GPU: split the rows, both sides compute
1219                // SIMULTANEOUSLY (same math, shared prescale).
1220                // GPU share: CMF_GPU_SPLIT (0..1, default 0.5).
1221                if *rows >= crate::gpu::min_rows()
1222                    && matches!(dtype, TensorDtype::Q8Row | TensorDtype::Q8_2f)
1223                    && gpu_lmhead_enabled()
1224                    && crate::gpu::enabled_here()
1225                {
1226                    // Runtime probe: alternate the hybrid against the
1227                    // pure-CPU matvec, keep whichever is faster HERE.
1228                    let t0 = std::time::Instant::now();
1229                    match crate::gpu::probe_arm(crate::gpu::OpClass::Matvec) {
1230                        crate::gpu::ProbeArm::Gpu => {}
1231                        crate::gpu::ProbeArm::CpuTimed => {
1232                            qmatvec(
1233                                self.quant_bytes(),
1234                                repack,
1235                                row_scale,
1236                                x,
1237                                col_field,
1238                                *dtype,
1239                                *rows,
1240                                *cols,
1241                                out,
1242                                pool,
1243                            );
1244                            crate::gpu::probe_record(
1245                                crate::gpu::OpClass::Matvec,
1246                                false,
1247                                t0.elapsed(),
1248                            );
1249                            return;
1250                        }
1251                        crate::gpu::ProbeArm::Cpu => {
1252                            qmatvec(
1253                                self.quant_bytes(),
1254                                repack,
1255                                row_scale,
1256                                x,
1257                                col_field,
1258                                *dtype,
1259                                *rows,
1260                                *cols,
1261                                out,
1262                                pool,
1263                            );
1264                            return;
1265                        }
1266                    }
1267                    let frac = gpu_split_frac();
1268                    let cpu_rows = ((*rows as f32) * (1.0 - frac)) as usize;
1269                    let (out_cpu, out_gpu) = out.split_at_mut(cpu_rows);
1270                    let bytes = self.quant_bytes();
1271                    let ok = std::thread::scope(|sc| {
1272                        let g = sc.spawn(|| {
1273                            crate::gpu::q8_matvec_range(
1274                                model,
1275                                *idx,
1276                                cpu_rows,
1277                                &row_scale[cpu_rows..],
1278                                &xs,
1279                                *rows - cpu_rows,
1280                                *cols,
1281                                out_gpu,
1282                            )
1283                        });
1284                        if cpu_rows > 0 {
1285                            // Repack prefix covers the full groups of the
1286                            // CPU half (the split starts at row 0).
1287                            let rep_cpu = if repack.is_empty() {
1288                                &[][..]
1289                            } else {
1290                                &repack[..(cpu_rows / 4) * 4 * *cols]
1291                            };
1292                            qmatvec(
1293                                &bytes[..cpu_rows * *cols],
1294                                rep_cpu,
1295                                &row_scale[..cpu_rows],
1296                                x,
1297                                col_field,
1298                                *dtype,
1299                                cpu_rows,
1300                                *cols,
1301                                out_cpu,
1302                                pool,
1303                            );
1304                        }
1305                        g.join().unwrap_or(false)
1306                    });
1307                    if ok {
1308                        crate::gpu::probe_record(crate::gpu::OpClass::Matvec, true, t0.elapsed());
1309                        return;
1310                    }
1311                    // GPU failed — CPU finishes its half (rows rebased —
1312                    // group offsets don't line up, mmap layout only).
1313                    qmatvec(
1314                        &bytes[cpu_rows * *cols..(*rows) * *cols],
1315                        &[],
1316                        &row_scale[cpu_rows..],
1317                        x,
1318                        col_field,
1319                        *dtype,
1320                        *rows - cpu_rows,
1321                        *cols,
1322                        out_gpu,
1323                        pool,
1324                    );
1325                    return;
1326                }
1327                qmatvec(
1328                    self.quant_bytes(),
1329                    repack,
1330                    row_scale,
1331                    x,
1332                    col_field,
1333                    *dtype,
1334                    *rows,
1335                    *cols,
1336                    out,
1337                    pool,
1338                );
1339            }
1340        }
1341    }
1342
1343    /// Fused two-input matvec (MTP verify pair): weights streamed once.
1344    pub fn matvec2(
1345        &self,
1346        x1: &[f32],
1347        x2: &[f32],
1348        o1: &mut [f32],
1349        o2: &mut [f32],
1350        pool: Option<&Pool>,
1351    ) {
1352        match self {
1353            Self::F32 { data, .. } => matvec_rows2(pool, data, x1, x2, o1, o2),
1354            Self::Mapped {
1355                dtype,
1356                rows,
1357                cols,
1358                row_scale,
1359                col_field,
1360                vbit_offsets,
1361                ..
1362            } => {
1363                if *dtype == TensorDtype::Q4Block {
1364                    q4matvec2(self.quant_bytes(), x1, x2, *rows, *cols, o1, o2, pool);
1365                    return;
1366                }
1367                if *dtype == TensorDtype::Q4Tiled {
1368                    q4t_matvec2(self.quant_bytes(), x1, x2, *rows, *cols, o1, o2, pool);
1369                    return;
1370                }
1371                if *dtype == TensorDtype::Q4TiledP {
1372                    q4tp_matvec2(self.quant_bytes(), x1, x2, *rows, *cols, o1, o2, pool);
1373                    return;
1374                }
1375                if *dtype == TensorDtype::Q2TiledP {
1376                    q2tp_matvec2(self.quant_bytes(), x1, x2, *rows, *cols, o1, o2, pool);
1377                    return;
1378                }
1379                if *dtype == TensorDtype::Q1 {
1380                    q1_matvec2(self.quant_bytes(), x1, x2, *rows, *cols, o1, o2, pool);
1381                    return;
1382                }
1383                if *dtype == TensorDtype::Q1T {
1384                    // Fused ternary pair: one row pass, the register
1385                    // unpack shared across both streams on ARM. (Q1T
1386                    // lacks a row_scale array — scales live inline in
1387                    // the tiles — so it must not fall through to the
1388                    // q8 qmatvec2 below.)
1389                    q1t_matvec2(self.quant_bytes(), x1, x2, *rows, *cols, o1, o2, pool);
1390                    return;
1391                }
1392                if matches!(dtype, TensorDtype::Vbit | TensorDtype::VbitRo) {
1393                    vbitmatvec2(
1394                        self.quant_bytes(),
1395                        vbit_offsets,
1396                        x1,
1397                        x2,
1398                        *rows,
1399                        *cols,
1400                        o1,
1401                        o2,
1402                        pool,
1403                    );
1404                    return;
1405                }
1406                qmatvec2(
1407                    self.quant_bytes(),
1408                    row_scale,
1409                    x1,
1410                    x2,
1411                    col_field,
1412                    *dtype,
1413                    *rows,
1414                    *cols,
1415                    o1,
1416                    o2,
1417                    pool,
1418                );
1419            }
1420        }
1421    }
1422}
1423
1424impl QTensor {
1425    /// Batched matvec (prefill-GEMM): xs — row-major [b, cols],
1426    /// out — row-major [b, rows]. Element-wise semantics are IDENTICAL
1427    /// to b matvec calls (same dot kernels in the same order); the win —
1428    /// the weight row streams from DRAM once per batch, not b times.
1429    /// `(model, index)` when this is a memory-mapped q4tp tensor — the
1430    /// identity a device-resident chain needs to hand `tp_matmat` the
1431    /// weight without going through this struct's own dispatch.
1432    pub fn q4tp_mapped(&self) -> Option<(&std::sync::Arc<CmfModel>, usize)> {
1433        match self {
1434            Self::Mapped {
1435                model, idx, dtype, ..
1436            } if *dtype == TensorDtype::Q4TiledP => Some((model, *idx)),
1437            _ => None,
1438        }
1439    }
1440
1441    pub fn matmat(&self, xs_all: &[f32], b: usize, out: &mut [f32], pool: Option<&Pool>) {
1442        let cols = self.cols();
1443        let rows = self.rows();
1444        debug_assert_eq!(xs_all.len(), b * cols);
1445        debug_assert_eq!(out.len(), b * rows);
1446        // GPTQ calibration: fold this layer's inputs into its Hessian. Only
1447        // Mapped tensors carry a directory name; the check is a relaxed
1448        // atomic load, free when not calibrating.
1449        if crate::gptq_capture::capturing() {
1450            if let Self::Mapped { model, idx, .. } = self {
1451                crate::gptq_capture::accumulate(&model.tensors[*idx].name, xs_all, b, cols);
1452            }
1453        }
1454        match self {
1455            Self::F32 { data, .. } => {
1456                let out_addr = SendMut(out.as_mut_ptr());
1457                let run = |start: usize, end: usize| {
1458                    for o in start..end {
1459                        let row = &data[o * cols..(o + 1) * cols];
1460                        for bi in 0..b {
1461                            let x = &xs_all[bi * cols..(bi + 1) * cols];
1462                            let mut acc = 0f32;
1463                            for j in 0..cols {
1464                                acc += row[j] * x[j];
1465                            }
1466                            unsafe { *out_addr.at(bi * rows + o) = acc };
1467                        }
1468                    }
1469                };
1470                dispatch_rows(pool, rows, &run);
1471            }
1472            Self::Mapped {
1473                dtype,
1474                row_scale,
1475                col_field,
1476                vbit_offsets,
1477                ..
1478            } => {
1479                if *dtype == TensorDtype::Q4Block {
1480                    q4matmat(self.quant_bytes(), xs_all, b, rows, cols, out, pool);
1481                    return;
1482                }
1483                if *dtype == TensorDtype::Q4TiledP {
1484                    // GPU batched q4tp GEMM (dequant + f32nt mul_mm on the
1485                    // device); the probe keeps whichever beats the CPU arm.
1486                    // Narrow (prompt-encode) and wide (DiT) batches probe
1487                    // as separate classes — the regimes have opposite
1488                    // winners and one shared verdict locked the wrong arm.
1489                    // Kill switch (gpu::mm_kill): one grossly slow GPU op
1490                    // (a fair-condition op is ≤~100 ms even at 1024px)
1491                    // means the device is contended by another process
1492                    // (e.g. a simulator) — verdicts are per-process, so
1493                    // without the bail the whole render crawls behind
1494                    // someone else's queue.
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                            // In-process A/B (`CMF_MM_AB=1`). Three
1508                            // wall-clock A/Bs on a shared stand disagreed
1509                            // with each other by 25% on the same change,
1510                            // because the machine drifts between processes
1511                            // and interleaving whole renders does not fix
1512                            // that. Here both arms run back to back on the
1513                            // SAME data inside one call, so whatever the
1514                            // machine is doing, it does to both — and the
1515                            // disagreement between their outputs falls out
1516                            // for free. Doubles the work; a diagnostic,
1517                            // not a mode.
1518                            if crate::mm_ab::on() {
1519                                let mut g = vec![0f32; b * rows];
1520                                let t = std::time::Instant::now();
1521                                let took = crate::gpu::q4tp_matmat(
1522                                    model, *idx, xs_all, b, rows, cols, &mut g,
1523                                );
1524                                let dg = t.elapsed();
1525                                let t = std::time::Instant::now();
1526                                q4tp_matmat(self.quant_bytes(), xs_all, b, rows, cols, out, pool);
1527                                let dc = t.elapsed();
1528                                crate::mm_ab::record(b, rows, cols, took, dg, dc, &g, out);
1529                                return;
1530                            }
1531                            let t0 = std::time::Instant::now();
1532                            // A cold call takes the device arm: its sample
1533                            // is discarded either way, and the upload is
1534                            // what the next step needs.
1535                            let resident = crate::gpu::weight_is_resident(model, *idx);
1536                            match crate::gpu::probe_arm_cold_prefers_gpu(class, resident) {
1537                                crate::gpu::ProbeArm::Gpu => {
1538                                    if crate::gpu::q4tp_matmat(
1539                                        model, *idx, xs_all, b, rows, cols, out,
1540                                    ) {
1541                                        let el = t0.elapsed();
1542                                        // Work-proportional budget: ~8× the
1543                                        // fair-device estimate (+20 ms slack).
1544                                        // An absolute cap missed the worst
1545                                        // case — contended ops sit at
1546                                        // 100–240 ms each and still bury a
1547                                        // render whose fair op is 3–9 ms.
1548                                        // Cold ops (first PSO build, buffer
1549                                        // alloc) are exempt: a one-off
1550                                        // ~50 ms compile is not contention.
1551                                        let flops = 2.0 * b as f64 * rows as f64 * cols as f64;
1552                                        let budget = std::time::Duration::from_secs_f64(
1553                                            flops / 1.5e12 * 8.0 + 0.020,
1554                                        );
1555                                        crate::gpu::mm_budget_check(
1556                                            "q4tp matmat",
1557                                            el,
1558                                            budget,
1559                                            crate::gpu::probe_was_cold() || !resident,
1560                                        );
1561                                        crate::gpu::probe_record(class, true, el);
1562                                        return;
1563                                    }
1564                                }
1565                                crate::gpu::ProbeArm::CpuTimed => {
1566                                    q4tp_matmat(
1567                                        self.quant_bytes(),
1568                                        xs_all,
1569                                        b,
1570                                        rows,
1571                                        cols,
1572                                        out,
1573                                        pool,
1574                                    );
1575                                    crate::gpu::probe_record(class, false, t0.elapsed());
1576                                    return;
1577                                }
1578                                crate::gpu::ProbeArm::Cpu => {}
1579                            }
1580                        }
1581                    }
1582                    q4tp_matmat(self.quant_bytes(), xs_all, b, rows, cols, out, pool);
1583                    return;
1584                }
1585                if *dtype == TensorDtype::Q2TiledP {
1586                    // Same device arm as q4tp, behind the same probe:
1587                    // the planes differ, the dispatch does not. Without
1588                    // this a q2tp file ran its widest projections on the
1589                    // host while the 4-bit one had the card, which is a
1590                    // codec paying for its size twice.
1591                    if b >= 32
1592                        && b * rows * cols >= 128_000_000
1593                        && cols % 32 == 0
1594                        && !crate::gpu::mm_killed()
1595                        && crate::gpu::enabled_here()
1596                    {
1597                        let class = if b >= 128 {
1598                            crate::gpu::OpClass::MatmatWide
1599                        } else {
1600                            crate::gpu::OpClass::Matmat
1601                        };
1602                        if let Self::Mapped { model, idx, .. } = self {
1603                            let t0 = std::time::Instant::now();
1604                            match crate::gpu::probe_arm(class) {
1605                                crate::gpu::ProbeArm::Gpu => {
1606                                    if crate::gpu::q2tp_matmat(
1607                                        model, *idx, xs_all, b, rows, cols, out,
1608                                    ) {
1609                                        crate::gpu::probe_record(class, true, t0.elapsed());
1610                                        return;
1611                                    }
1612                                }
1613                                crate::gpu::ProbeArm::CpuTimed => {
1614                                    q2tp_matmat(
1615                                        self.quant_bytes(),
1616                                        xs_all,
1617                                        b,
1618                                        rows,
1619                                        cols,
1620                                        out,
1621                                        pool,
1622                                    );
1623                                    crate::gpu::probe_record(class, false, t0.elapsed());
1624                                    return;
1625                                }
1626                                crate::gpu::ProbeArm::Cpu => {}
1627                            }
1628                        }
1629                    }
1630                    // Without a host arm a q2tp tensor falls through to
1631                    // the q8 fallback, which reads it at one BYTE per
1632                    // weight — a 2x overrun that killed pool workers
1633                    // mid-prefill while the dispatcher waited forever.
1634                    q2tp_matmat(self.quant_bytes(), xs_all, b, rows, cols, out, pool);
1635                    return;
1636                }
1637                if *dtype == TensorDtype::Q4Tiled {
1638                    // GPU batched q4t GEMM (dequant + f32nt mul_mm on the
1639                    // device); the probe keeps whichever beats the CPU arm.
1640                    // Narrow (prompt-encode) and wide (DiT) batches probe
1641                    // as separate classes — the regimes have opposite
1642                    // winners and one shared verdict locked the wrong arm.
1643                    // Kill switch (gpu::mm_kill): one grossly slow GPU op
1644                    // (a fair-condition op is ≤~100 ms even at 1024px)
1645                    // means the device is contended by another process
1646                    // (e.g. a simulator) — verdicts are per-process, so
1647                    // without the bail the whole render crawls behind
1648                    // someone else's queue.
1649                    if b >= 32
1650                        && b * rows * cols >= 128_000_000
1651                        && cols % 32 == 0
1652                        && !crate::gpu::mm_killed()
1653                        && crate::gpu::enabled_here()
1654                    {
1655                        let class = if b >= 128 {
1656                            crate::gpu::OpClass::MatmatWide
1657                        } else {
1658                            crate::gpu::OpClass::Matmat
1659                        };
1660                        if let Self::Mapped { model, idx, .. } = self {
1661                            let t0 = std::time::Instant::now();
1662                            match crate::gpu::probe_arm(class) {
1663                                crate::gpu::ProbeArm::Gpu => {
1664                                    if crate::gpu::q4t_matmat(
1665                                        model, *idx, xs_all, b, rows, cols, out,
1666                                    ) {
1667                                        let el = t0.elapsed();
1668                                        // Work-proportional budget: ~8× the
1669                                        // fair-device estimate (+20 ms slack).
1670                                        // An absolute cap missed the worst
1671                                        // case — contended ops sit at
1672                                        // 100–240 ms each and still bury a
1673                                        // render whose fair op is 3–9 ms.
1674                                        // Cold ops (first PSO build, buffer
1675                                        // alloc) are exempt: a one-off
1676                                        // ~50 ms compile is not contention.
1677                                        let flops = 2.0 * b as f64 * rows as f64 * cols as f64;
1678                                        let budget = std::time::Duration::from_secs_f64(
1679                                            flops / 1.5e12 * 8.0 + 0.020,
1680                                        );
1681                                        crate::gpu::mm_budget_check(
1682                                            "q4t matmat",
1683                                            el,
1684                                            budget,
1685                                            crate::gpu::probe_was_cold(),
1686                                        );
1687                                        crate::gpu::probe_record(class, true, el);
1688                                        return;
1689                                    }
1690                                }
1691                                crate::gpu::ProbeArm::CpuTimed => {
1692                                    q4t_matmat(
1693                                        self.quant_bytes(),
1694                                        xs_all,
1695                                        b,
1696                                        rows,
1697                                        cols,
1698                                        out,
1699                                        pool,
1700                                    );
1701                                    crate::gpu::probe_record(class, false, t0.elapsed());
1702                                    return;
1703                                }
1704                                crate::gpu::ProbeArm::Cpu => {}
1705                            }
1706                        }
1707                    }
1708                    q4t_matmat(self.quant_bytes(), xs_all, b, rows, cols, out, pool);
1709                    return;
1710                }
1711                if *dtype == TensorDtype::Q1 {
1712                    // GPU batched q1 GEMM for wide prefill (q1_mul_mm on the
1713                    // device); the probe keeps whichever beats the CPU matmat.
1714                    if b >= 32
1715                        && b * rows * cols >= 128_000_000
1716                        && cols % 64 == 0
1717                        && crate::gpu::enabled_here()
1718                    {
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::q1_matmat(
1724                                        model, *idx, xs_all, b, rows, cols, out,
1725                                    ) {
1726                                        crate::gpu::probe_record(
1727                                            crate::gpu::OpClass::Matmat,
1728                                            true,
1729                                            t0.elapsed(),
1730                                        );
1731                                        return;
1732                                    }
1733                                }
1734                                crate::gpu::ProbeArm::CpuTimed => {
1735                                    q1_matmat(self.quant_bytes(), xs_all, b, rows, cols, out, pool);
1736                                    crate::gpu::probe_record(
1737                                        crate::gpu::OpClass::Matmat,
1738                                        false,
1739                                        t0.elapsed(),
1740                                    );
1741                                    return;
1742                                }
1743                                crate::gpu::ProbeArm::Cpu => {}
1744                            }
1745                        }
1746                    }
1747                    q1_matmat(self.quant_bytes(), xs_all, b, rows, cols, out, pool);
1748                    return;
1749                }
1750                if *dtype == TensorDtype::Q1T {
1751                    // GPU batched GEMM for wide prefill (base + overlay on the
1752                    // device); probe keeps the winner vs the CPU matmat.
1753                    if b >= 32 && b * rows * cols >= 128_000_000 && crate::gpu::enabled_here() {
1754                        if let Self::Mapped { model, idx, .. } = self {
1755                            let t0 = std::time::Instant::now();
1756                            match crate::gpu::probe_arm(crate::gpu::OpClass::Matmat) {
1757                                crate::gpu::ProbeArm::Gpu => {
1758                                    if crate::gpu::q1t_matmat(
1759                                        model, *idx, xs_all, b, rows, cols, out,
1760                                    ) {
1761                                        crate::gpu::probe_record(
1762                                            crate::gpu::OpClass::Matmat,
1763                                            true,
1764                                            t0.elapsed(),
1765                                        );
1766                                        return;
1767                                    }
1768                                }
1769                                crate::gpu::ProbeArm::CpuTimed => {
1770                                    q1t_matmat(
1771                                        self.quant_bytes(),
1772                                        xs_all,
1773                                        b,
1774                                        rows,
1775                                        cols,
1776                                        out,
1777                                        pool,
1778                                    );
1779                                    crate::gpu::probe_record(
1780                                        crate::gpu::OpClass::Matmat,
1781                                        false,
1782                                        t0.elapsed(),
1783                                    );
1784                                    return;
1785                                }
1786                                crate::gpu::ProbeArm::Cpu => {}
1787                            }
1788                        }
1789                    }
1790                    q1t_matmat(self.quant_bytes(), xs_all, b, rows, cols, out, pool);
1791                    return;
1792                }
1793                if matches!(dtype, TensorDtype::Vbit | TensorDtype::VbitRo) {
1794                    vbitmatmat(
1795                        self.quant_bytes(),
1796                        vbit_offsets,
1797                        xs_all,
1798                        b,
1799                        rows,
1800                        cols,
1801                        out,
1802                        pool,
1803                    );
1804                    return;
1805                }
1806                let pre: Vec<std::borrow::Cow<'_, [f32]>> = (0..b)
1807                    .map(|bi| prescale(&xs_all[bi * cols..(bi + 1) * cols], col_field, *dtype))
1808                    .collect();
1809                // D5: large prefill-batch GEMMs — on the GPU (threshold by
1810                // work volume: submission carries b×rows×cols MACs).
1811                // Runtime probe: the naive GEMM shader + sync readback
1812                // lose to the CPU GEMM on slow driver stacks — alternate
1813                // both arms and keep the winner.
1814                if b >= 8 && b * rows * cols >= 128_000_000 && crate::gpu::enabled_here() {
1815                    if let Self::Mapped { model, idx, .. } = self {
1816                        let t0 = std::time::Instant::now();
1817                        match crate::gpu::probe_arm(crate::gpu::OpClass::Matmat) {
1818                            crate::gpu::ProbeArm::Gpu
1819                                if crate::gpu::probe_deciding(crate::gpu::OpClass::Matmat)
1820                                    && !crate::gpu::q8_resident_or_upload(model, *idx) =>
1821                            {
1822                                // Cold weights during probing: the upload
1823                                // has started, the count runs on the CPU —
1824                                // the GPU arm samples on the next touch.
1825                                let q = self.quant_bytes();
1826                                qmatmat(q, row_scale, &pre, rows, cols, out, pool);
1827                                return;
1828                            }
1829                            crate::gpu::ProbeArm::Gpu => {
1830                                let flat: Vec<f32> =
1831                                    pre.iter().flat_map(|v| v.iter().copied()).collect();
1832                                if crate::gpu::q8_matmat(
1833                                    model, *idx, row_scale, &flat, b, rows, cols, out,
1834                                ) {
1835                                    crate::gpu::probe_record(
1836                                        crate::gpu::OpClass::Matmat,
1837                                        true,
1838                                        t0.elapsed(),
1839                                    );
1840                                    return;
1841                                }
1842                            }
1843                            crate::gpu::ProbeArm::CpuTimed => {
1844                                let q = self.quant_bytes();
1845                                qmatmat(q, row_scale, &pre, rows, cols, out, pool);
1846                                crate::gpu::probe_record(
1847                                    crate::gpu::OpClass::Matmat,
1848                                    false,
1849                                    t0.elapsed(),
1850                                );
1851                                return;
1852                            }
1853                            crate::gpu::ProbeArm::Cpu => {}
1854                        }
1855                    }
1856                }
1857                let q = self.quant_bytes();
1858                qmatmat(q, row_scale, &pre, rows, cols, out, pool);
1859            }
1860        }
1861    }
1862}
1863
1864impl QTensor {
1865    /// The device GEMM this tensor would take, run once on the caller's
1866    /// data — the startup parity probe's arm, and the one place that knows
1867    /// which entry point each codec has.
1868    ///
1869    /// It exists because the probe used to look for a `q4tp` weight by
1870    /// name AND dtype, and a container packed any other way was declared
1871    /// "host path" for the whole render even though its codec had a device
1872    /// GEMM of its own. A gate that only recognizes one codec is a gate
1873    /// that silently downgrades every other one.
1874    pub fn device_matmat(&self, xs: &[f32], b: usize, out: &mut [f32]) -> bool {
1875        let (rows, cols) = (self.rows(), self.cols());
1876        let Self::Mapped {
1877            model,
1878            idx,
1879            dtype,
1880            row_scale,
1881            col_field,
1882            ..
1883        } = self
1884        else {
1885            return false;
1886        };
1887        match *dtype {
1888            TensorDtype::Q4TiledP => crate::gpu::q4tp_matmat(model, *idx, xs, b, rows, cols, out),
1889            // The two-field codec folds its column field into the
1890            // activation, which leaves a plain per-row int8 GEMM — the
1891            // same kernel `q8_row` uses, on both backends.
1892            TensorDtype::Q8Row | TensorDtype::Q8_2f => {
1893                // The field belongs to the weight; only a backend that cannot
1894                // apply it there makes a scaled copy of the activation.
1895                if *dtype == TensorDtype::Q8_2f
1896                    && std::env::var("CMF_Q8_2F_DEV").as_deref() != Ok("0")
1897                    && crate::gpu::q8_matmat_2f(
1898                        model, *idx, row_scale, col_field, xs, b, rows, cols, out,
1899                    )
1900                {
1901                    return true;
1902                }
1903                let flat: Vec<f32> = (0..b)
1904                    .flat_map(|bi| {
1905                        prescale(&xs[bi * cols..(bi + 1) * cols], col_field, *dtype).into_owned()
1906                    })
1907                    .collect();
1908                crate::gpu::q8_matmat(model, *idx, row_scale, &flat, b, rows, cols, out)
1909            }
1910            _ => false,
1911        }
1912    }
1913
1914    /// Multi-matrix job (roadmap §3 P0): N tensors sharing one input
1915    /// run under a SINGLE pool dispatch — QKV or gate+up cost one
1916    /// barrier instead of N. Per-row math is the exact same kernel as
1917    /// `matvec` (bit-identical outputs); only the dispatch is fused.
1918    /// Falls back to N sequential matvecs when the set is not a uniform
1919    /// q8-family/F32 group or there is no pool.
1920    pub fn matvec_many<const N: usize>(
1921        ts: [&QTensor; N],
1922        x: &[f32],
1923        mut outs: [&mut [f32]; N],
1924        pool: Option<&Pool>,
1925    ) {
1926        let total_rows: usize = ts.iter().map(|t| t.rows()).sum();
1927        let uniform_q8 = ts.iter().all(|t| {
1928            matches!(
1929                t,
1930                Self::Mapped {
1931                    dtype: TensorDtype::Q8Row | TensorDtype::Q8_2f,
1932                    ..
1933                }
1934            )
1935        });
1936        let uniform_f32 = ts.iter().all(|t| matches!(t, Self::F32 { .. }));
1937        let uniform_q4 = ts.iter().all(|t| {
1938            matches!(
1939                t,
1940                Self::Mapped {
1941                    dtype: TensorDtype::Q4Block,
1942                    ..
1943                }
1944            )
1945        });
1946        let uniform_vbit = ts.iter().all(|t| {
1947            matches!(
1948                t,
1949                Self::Mapped {
1950                    dtype: TensorDtype::Vbit | TensorDtype::VbitRo,
1951                    ..
1952                }
1953            )
1954        });
1955        let uniform_q1 = ts.iter().all(|t| {
1956            matches!(
1957                t,
1958                Self::Mapped {
1959                    dtype: TensorDtype::Q1,
1960                    ..
1961                }
1962            )
1963        });
1964        let uniform_q1t = ts.iter().all(|t| {
1965            matches!(
1966                t,
1967                Self::Mapped {
1968                    dtype: TensorDtype::Q1T,
1969                    ..
1970                }
1971            )
1972        });
1973        // q4tp is the skeleton dtype of the big MoE files, and without an arm
1974        // here every projection that shares an input paid its own pool
1975        // barrier: DeepSeek-V4's attention step alone hands this function
1976        // wq_a, wkv and both compressors' pairs off the same hidden state.
1977        let uniform_q4tp = ts.iter().all(|t| {
1978            matches!(
1979                t,
1980                Self::Mapped {
1981                    dtype: TensorDtype::Q4TiledP,
1982                    ..
1983                }
1984            )
1985        }) && ts
1986            .iter()
1987            .all(|t| t.cols() == ts[0].cols() && t.cols() % GROUP_SIZE == 0);
1988        let Some(pool) = pool else {
1989            for (t, o) in ts.iter().zip(outs.iter_mut()) {
1990                t.matvec(x, o, None);
1991            }
1992            return;
1993        };
1994        if total_rows < 256
1995            || !(uniform_q8
1996                || uniform_f32
1997                || uniform_q4
1998                || uniform_vbit
1999                || uniform_q1
2000                || uniform_q1t
2001                || uniform_q4tp)
2002        {
2003            for (t, o) in ts.iter().zip(outs.iter_mut()) {
2004                t.matvec(x, o, Some(pool));
2005            }
2006            return;
2007        }
2008
2009        if uniform_q4tp {
2010            // Every tensor's rows laid end to end in one virtual row space,
2011            // so the whole set is ONE dispatch. The per-row body is the
2012            // `q4tp_matvec` arm verbatim — same activation split, same
2013            // accumulation order — so the outputs are bit-identical to the
2014            // sequential calls this replaces.
2015            let cols = ts[0].cols();
2016            let gpr = cols / GROUP_SIZE;
2017            let views: Vec<Q4tpView> = ts
2018                .iter()
2019                .map(|t| Q4tpView::new(t.quant_bytes(), t.rows(), cols))
2020                .collect();
2021            let rows_of: Vec<usize> = ts.iter().map(|t| t.rows()).collect();
2022            let outs_addr: [SendMut; N] = std::array::from_fn(|i| SendMut(outs[i].as_mut_ptr()));
2023            // flat index -> (which tensor, which of its rows)
2024            let locate = |flat: usize| -> (usize, usize) {
2025                let mut acc = 0;
2026                for (i, &r) in rows_of.iter().enumerate() {
2027                    if flat < acc + r {
2028                        return (i, flat - acc);
2029                    }
2030                    acc += r;
2031                }
2032                (rows_of.len() - 1, 0)
2033            };
2034            let (views, outs_addr) = (&views, &outs_addr);
2035            if a8w8_enabled() {
2036                let act = split_act(x);
2037                let act = &act;
2038                let run = |start: usize, end: usize| {
2039                    let mut sc = vec![0f32; gpr];
2040                    for flat in start..end {
2041                        let (t, r) = locate(flat);
2042                        let v = &views[t];
2043                        v.scales_into(r, gpr, &mut sc);
2044                        let mut acc = dot_q4tp_row_i8(v.nib, r, gpr, &act.xq, &sc) * act.sx;
2045                        for &(j, xv) in &act.outliers {
2046                            let (w, s) = q4tp_outlier(v.nib, r, gpr, j, &sc);
2047                            acc += w * s * xv;
2048                        }
2049                        // SAFETY: one worker owns each (tensor, row) pair.
2050                        unsafe { *outs_addr[t].at(r) = acc };
2051                    }
2052                };
2053                pool.run_rows(total_rows, &run);
2054            } else {
2055                let run = |start: usize, end: usize| {
2056                    let mut sc = vec![0f32; gpr];
2057                    for flat in start..end {
2058                        let (t, r) = locate(flat);
2059                        let v = &views[t];
2060                        v.scales_into(r, gpr, &mut sc);
2061                        // SAFETY: one worker owns each (tensor, row) pair.
2062                        unsafe { *outs_addr[t].at(r) = q4tp_row_exact(v.nib, r, gpr, x, &sc) };
2063                    }
2064                };
2065                pool.run_rows(total_rows, &run);
2066            }
2067            return;
2068        }
2069
2070        if uniform_q1 {
2071            // One shared activation split + group sums (q1 has no col
2072            // field; the same input feeds every tensor).
2073            let outs_addr: [SendMut; N] = std::array::from_fn(|i| SendMut(outs[i].as_mut_ptr()));
2074            if a8w8_enabled() {
2075                let act = split_act(x);
2076                let gsum = q1_group_sums(&act.xq, ts[0].cols() / GROUP_SIZE);
2077                let (act, gsum) = (&act, &gsum);
2078                let closures: [_; N] = std::array::from_fn(|i| {
2079                    let (bytes, gpr, out) =
2080                        (ts[i].quant_bytes(), ts[i].cols() / GROUP_SIZE, outs_addr[i]);
2081                    move |s: usize, e: usize| q1_range_a8w8(bytes, gpr, act, gsum, out, s, e)
2082                });
2083                let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
2084                    std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
2085                pool.run_many(&parts);
2086            } else {
2087                let closures: [_; N] = std::array::from_fn(|i| {
2088                    let (bytes, gpr, out) =
2089                        (ts[i].quant_bytes(), ts[i].cols() / GROUP_SIZE, outs_addr[i]);
2090                    move |s: usize, e: usize| q1_range_f32(bytes, gpr, x, out, s, e)
2091                });
2092                let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
2093                    std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
2094                pool.run_many(&parts);
2095            }
2096            return;
2097        }
2098
2099        if uniform_q1t {
2100            // Q1T batched: one shared activation split + overlay decode,
2101            // all tensors' rows in ONE pool dispatch (saves N−1 dispatches
2102            // and N−1 redundant split_act calls per layer).
2103            let outs_addr: [SendMut; N] = std::array::from_fn(|i| SendMut(outs[i].as_mut_ptr()));
2104            const TILE: usize = cortiq_core::quant::Q1T_TILE;
2105            if a8w8_enabled() {
2106                let act = split_act(x);
2107                let act = &act;
2108                let x_ref = x;
2109                let closures: [_; N] = std::array::from_fn(|i| {
2110                    let bytes = ts[i].quant_bytes();
2111                    let (rows, cols) = (ts[i].rows(), ts[i].cols());
2112                    let gpr = cols / GROUP_SIZE;
2113                    let (rp_off, ent_off, has_ov) = q1t_overlay(bytes, rows * gpr * TILE, rows);
2114                    let out = outs_addr[i];
2115                    move |s: usize, e: usize| {
2116                        q1t_range_a8w8(bytes, gpr, rp_off, ent_off, has_ov, act, x_ref, out, s, e)
2117                    }
2118                });
2119                let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
2120                    std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
2121                pool.run_many(&parts);
2122            } else {
2123                let x_ref = x;
2124                let closures: [_; N] = std::array::from_fn(|i| {
2125                    let bytes = ts[i].quant_bytes();
2126                    let (rows, cols) = (ts[i].rows(), ts[i].cols());
2127                    let gpr = cols / GROUP_SIZE;
2128                    let (rp_off, ent_off, has_ov) = q1t_overlay(bytes, rows * gpr * TILE, rows);
2129                    let out = outs_addr[i];
2130                    move |s: usize, e: usize| {
2131                        q1t_range_f32_batch(bytes, gpr, rp_off, ent_off, has_ov, x_ref, out, s, e)
2132                    }
2133                });
2134                let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
2135                    std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
2136                pool.run_many(&parts);
2137            }
2138            return;
2139        }
2140
2141        if uniform_q4 || uniform_vbit {
2142            let outs_addr: [SendMut; N] = std::array::from_fn(|i| SendMut(outs[i].as_mut_ptr()));
2143            // q4/vbit share one activation split — no per-tensor col field.
2144            if a8w8_enabled() {
2145                let act = split_act(x);
2146                let act = &act;
2147                if uniform_q4 {
2148                    let closures: [_; N] = std::array::from_fn(|i| {
2149                        let (packed, scales) =
2150                            q4_split(ts[i].quant_bytes(), ts[i].rows(), ts[i].cols());
2151                        let (gpr, cols, out) =
2152                            (ts[i].cols() / GROUP_SIZE, ts[i].cols(), outs_addr[i]);
2153                        move |s: usize, e: usize| {
2154                            q4_range_a8w8(packed, scales, gpr, cols, act, out, s, e)
2155                        }
2156                    });
2157                    let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
2158                        std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
2159                    pool.run_many(&parts);
2160                } else {
2161                    let closures: [_; N] = std::array::from_fn(|i| {
2162                        let Self::Mapped { vbit_offsets, .. } = ts[i] else {
2163                            unreachable!()
2164                        };
2165                        let (bytes, rows, cols, out) = (
2166                            ts[i].quant_bytes(),
2167                            ts[i].rows(),
2168                            ts[i].cols(),
2169                            outs_addr[i],
2170                        );
2171                        move |s: usize, e: usize| {
2172                            vbit_range_a8w8(bytes, vbit_offsets, x, act, rows, cols, out, s, e)
2173                        }
2174                    });
2175                    let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
2176                        std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
2177                    pool.run_many(&parts);
2178                }
2179                return;
2180            }
2181            if uniform_q4 {
2182                let closures: [_; N] = std::array::from_fn(|i| {
2183                    let (packed, scales) =
2184                        q4_split(ts[i].quant_bytes(), ts[i].rows(), ts[i].cols());
2185                    let (gpr, out) = (ts[i].cols() / GROUP_SIZE, outs_addr[i]);
2186                    move |s: usize, e: usize| q4_range_f32(packed, scales, gpr, x, out, s, e)
2187                });
2188                let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
2189                    std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
2190                pool.run_many(&parts);
2191            } else {
2192                let closures: [_; N] = std::array::from_fn(|i| {
2193                    let Self::Mapped { vbit_offsets, .. } = ts[i] else {
2194                        unreachable!()
2195                    };
2196                    let (bytes, rows, cols, out) = (
2197                        ts[i].quant_bytes(),
2198                        ts[i].rows(),
2199                        ts[i].cols(),
2200                        outs_addr[i],
2201                    );
2202                    move |s: usize, e: usize| {
2203                        vbit_range_f32(bytes, vbit_offsets, x, rows, cols, out, s, e)
2204                    }
2205                });
2206                let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
2207                    std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
2208                pool.run_many(&parts);
2209            }
2210            return;
2211        }
2212
2213        if uniform_f32 {
2214            let outs_addr: [SendMut; N] = std::array::from_fn(|i| SendMut(outs[i].as_mut_ptr()));
2215            let closures: [_; N] = std::array::from_fn(|i| {
2216                let Self::F32 { data, cols, .. } = ts[i] else {
2217                    unreachable!()
2218                };
2219                let out = outs_addr[i];
2220                move |start: usize, end: usize| {
2221                    for o in start..end {
2222                        let row = &data[o * cols..(o + 1) * cols];
2223                        let mut sum = 0.0f32;
2224                        for j in 0..*cols {
2225                            sum += row[j] * x[j];
2226                        }
2227                        // SAFETY: disjoint (tensor, row) cells per worker.
2228                        unsafe { *out.at(o) = sum };
2229                    }
2230                }
2231            });
2232            let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
2233                std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
2234            pool.run_many(&parts);
2235            return;
2236        }
2237
2238        // Uniform q8-family: per-tensor prescale (q8_2f col fields
2239        // differ per tensor) + the shared range kernels.
2240        struct Ctx<'a> {
2241            bytes: &'a [u8],
2242            #[cfg_attr(not(target_arch = "aarch64"), allow(dead_code))]
2243            rep: &'a [u8],
2244            row_scale: &'a [f32],
2245            cols: usize,
2246            xs: std::borrow::Cow<'a, [f32]>,
2247        }
2248        let ctxs: [Ctx<'_>; N] = std::array::from_fn(|i| {
2249            let Self::Mapped {
2250                dtype,
2251                cols,
2252                row_scale,
2253                col_field,
2254                repack,
2255                ..
2256            } = ts[i]
2257            else {
2258                unreachable!()
2259            };
2260            Ctx {
2261                bytes: ts[i].quant_bytes(),
2262                rep: repack,
2263                row_scale,
2264                cols: *cols,
2265                xs: prescale(x, col_field, *dtype),
2266            }
2267        });
2268        let outs_addr: [SendMut; N] = std::array::from_fn(|i| SendMut(outs[i].as_mut_ptr()));
2269        #[cfg(target_arch = "aarch64")]
2270        if sdot_enabled() {
2271            let acts: [SplitAct; N] = std::array::from_fn(|i| split_act(&ctxs[i].xs));
2272            let closures: [_; N] = std::array::from_fn(|i| {
2273                let (c, act, out) = (&ctxs[i], &acts[i], outs_addr[i]);
2274                move |start: usize, end: usize| {
2275                    q8_range_sdot(c.bytes, c.rep, c.row_scale, act, c.cols, out, start, end)
2276                }
2277            });
2278            let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
2279                std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
2280            pool.run_many(&parts);
2281            return;
2282        }
2283        #[cfg(target_arch = "x86_64")]
2284        if avx2_a8w8_enabled() {
2285            let acts: [SplitAct; N] = std::array::from_fn(|i| split_act(&ctxs[i].xs));
2286            let closures: [_; N] = std::array::from_fn(|i| {
2287                let (c, act, out) = (&ctxs[i], &acts[i], outs_addr[i]);
2288                move |start: usize, end: usize| {
2289                    q8_range_avx2(c.bytes, c.row_scale, act, c.cols, out, start, end)
2290                }
2291            });
2292            let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
2293                std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
2294            pool.run_many(&parts);
2295            return;
2296        }
2297        let closures: [_; N] = std::array::from_fn(|i| {
2298            let (c, out) = (&ctxs[i], outs_addr[i]);
2299            move |start: usize, end: usize| {
2300                q8_range_f32(c.bytes, c.row_scale, &c.xs, c.cols, out, start, end)
2301            }
2302        });
2303        let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
2304            std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
2305        pool.run_many(&parts);
2306    }
2307}
2308
2309impl QTensor {
2310    /// Pair-input multi-matrix job: N tensors × 2 shared inputs under a
2311    /// single pool dispatch — the MTP/pair decode path publishes one job
2312    /// for Q/K/V (and one for gate+up) instead of one per tensor.
2313    /// Per-row math is exactly `matvec2`'s kernels; bit-identical.
2314    #[allow(clippy::needless_range_loop)]
2315    pub fn matvec2_many<const N: usize>(
2316        ts: [&QTensor; N],
2317        x1: &[f32],
2318        x2: &[f32],
2319        mut o1s: [&mut [f32]; N],
2320        mut o2s: [&mut [f32]; N],
2321        pool: Option<&Pool>,
2322    ) {
2323        let total_rows: usize = ts.iter().map(|t| t.rows()).sum();
2324        let uniform_q8 = ts.iter().all(|t| {
2325            matches!(
2326                t,
2327                Self::Mapped {
2328                    dtype: TensorDtype::Q8Row | TensorDtype::Q8_2f,
2329                    ..
2330                }
2331            )
2332        });
2333        let uniform_f32 = ts.iter().all(|t| matches!(t, Self::F32 { .. }));
2334        let uniform_q4 = ts.iter().all(|t| {
2335            matches!(
2336                t,
2337                Self::Mapped {
2338                    dtype: TensorDtype::Q4Block,
2339                    ..
2340                }
2341            )
2342        });
2343        let uniform_vbit = ts.iter().all(|t| {
2344            matches!(
2345                t,
2346                Self::Mapped {
2347                    dtype: TensorDtype::Vbit | TensorDtype::VbitRo,
2348                    ..
2349                }
2350            )
2351        });
2352        let fusable = pool.is_some()
2353            && total_rows >= 256
2354            && (uniform_q8 || uniform_f32 || uniform_q4 || uniform_vbit);
2355        if !fusable {
2356            for i in 0..N {
2357                ts[i].matvec2(x1, x2, o1s[i], o2s[i], pool);
2358            }
2359            return;
2360        }
2361        let pool = pool.unwrap();
2362
2363        if uniform_q4 || uniform_vbit {
2364            let p1: [SendMut; N] = std::array::from_fn(|i| SendMut(o1s[i].as_mut_ptr()));
2365            let p2: [SendMut; N] = std::array::from_fn(|i| SendMut(o2s[i].as_mut_ptr()));
2366            // q4/vbit share activation splits — no per-tensor col field.
2367            if a8w8_enabled() {
2368                let a1 = split_act(x1);
2369                let a2 = split_act(x2);
2370                let (a1, a2) = (&a1, &a2);
2371                if uniform_q4 {
2372                    let closures: [_; N] = std::array::from_fn(|i| {
2373                        let (packed, scales) =
2374                            q4_split(ts[i].quant_bytes(), ts[i].rows(), ts[i].cols());
2375                        let (gpr, cols, o1, o2) =
2376                            (ts[i].cols() / GROUP_SIZE, ts[i].cols(), p1[i], p2[i]);
2377                        move |s: usize, e: usize| {
2378                            q4_range2_a8w8(packed, scales, gpr, cols, a1, a2, o1, o2, s, e)
2379                        }
2380                    });
2381                    let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
2382                        std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
2383                    pool.run_many(&parts);
2384                } else {
2385                    let closures: [_; N] = std::array::from_fn(|i| {
2386                        let Self::Mapped { vbit_offsets, .. } = ts[i] else {
2387                            unreachable!()
2388                        };
2389                        let (bytes, rows, cols, o1, o2) = (
2390                            ts[i].quant_bytes(),
2391                            ts[i].rows(),
2392                            ts[i].cols(),
2393                            p1[i],
2394                            p2[i],
2395                        );
2396                        move |s: usize, e: usize| {
2397                            vbit_range2_a8w8(
2398                                bytes,
2399                                vbit_offsets,
2400                                x1,
2401                                x2,
2402                                a1,
2403                                a2,
2404                                rows,
2405                                cols,
2406                                o1,
2407                                o2,
2408                                s,
2409                                e,
2410                            )
2411                        }
2412                    });
2413                    let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
2414                        std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
2415                    pool.run_many(&parts);
2416                }
2417                return;
2418            }
2419            if uniform_q4 {
2420                let closures: [_; N] = std::array::from_fn(|i| {
2421                    let (packed, scales) =
2422                        q4_split(ts[i].quant_bytes(), ts[i].rows(), ts[i].cols());
2423                    let (gpr, o1, o2) = (ts[i].cols() / GROUP_SIZE, p1[i], p2[i]);
2424                    move |s: usize, e: usize| {
2425                        q4_range2_f32(packed, scales, gpr, x1, x2, o1, o2, s, e)
2426                    }
2427                });
2428                let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
2429                    std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
2430                pool.run_many(&parts);
2431            } else {
2432                let closures: [_; N] = std::array::from_fn(|i| {
2433                    let Self::Mapped { vbit_offsets, .. } = ts[i] else {
2434                        unreachable!()
2435                    };
2436                    let (bytes, rows, cols, o1, o2) = (
2437                        ts[i].quant_bytes(),
2438                        ts[i].rows(),
2439                        ts[i].cols(),
2440                        p1[i],
2441                        p2[i],
2442                    );
2443                    move |s: usize, e: usize| {
2444                        vbit_range2_f32(bytes, vbit_offsets, x1, x2, rows, cols, o1, o2, s, e)
2445                    }
2446                });
2447                let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
2448                    std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
2449                pool.run_many(&parts);
2450            }
2451            return;
2452        }
2453
2454        if uniform_f32 {
2455            let p1: [SendMut; N] = std::array::from_fn(|i| SendMut(o1s[i].as_mut_ptr()));
2456            let p2: [SendMut; N] = std::array::from_fn(|i| SendMut(o2s[i].as_mut_ptr()));
2457            let closures: [_; N] = std::array::from_fn(|i| {
2458                let Self::F32 { data, cols, .. } = ts[i] else {
2459                    unreachable!()
2460                };
2461                let (o1, o2) = (p1[i], p2[i]);
2462                move |start: usize, end: usize| {
2463                    for o in start..end {
2464                        let row = &data[o * cols..(o + 1) * cols];
2465                        let (mut s1, mut s2) = (0.0f32, 0.0f32);
2466                        for j in 0..*cols {
2467                            s1 += row[j] * x1[j];
2468                            s2 += row[j] * x2[j];
2469                        }
2470                        // SAFETY: disjoint (tensor, row) cells per worker.
2471                        unsafe {
2472                            *o1.at(o) = s1;
2473                            *o2.at(o) = s2;
2474                        }
2475                    }
2476                }
2477            });
2478            let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
2479                std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
2480            pool.run_many(&parts);
2481            return;
2482        }
2483
2484        struct Ctx<'a> {
2485            bytes: &'a [u8],
2486            row_scale: &'a [f32],
2487            cols: usize,
2488            xs1: std::borrow::Cow<'a, [f32]>,
2489            xs2: std::borrow::Cow<'a, [f32]>,
2490        }
2491        let ctxs: [Ctx<'_>; N] = std::array::from_fn(|i| {
2492            let Self::Mapped {
2493                dtype,
2494                cols,
2495                row_scale,
2496                col_field,
2497                ..
2498            } = ts[i]
2499            else {
2500                unreachable!()
2501            };
2502            Ctx {
2503                bytes: ts[i].quant_bytes(),
2504                row_scale,
2505                cols: *cols,
2506                xs1: prescale(x1, col_field, *dtype),
2507                xs2: prescale(x2, col_field, *dtype),
2508            }
2509        });
2510        let p1: [SendMut; N] = std::array::from_fn(|i| SendMut(o1s[i].as_mut_ptr()));
2511        let p2: [SendMut; N] = std::array::from_fn(|i| SendMut(o2s[i].as_mut_ptr()));
2512        #[cfg(target_arch = "aarch64")]
2513        if sdot_enabled() {
2514            let acts: [(SplitAct, SplitAct); N] =
2515                std::array::from_fn(|i| (split_act(&ctxs[i].xs1), split_act(&ctxs[i].xs2)));
2516            let closures: [_; N] = std::array::from_fn(|i| {
2517                let (c, a, o1, o2) = (&ctxs[i], &acts[i], p1[i], p2[i]);
2518                move |start: usize, end: usize| {
2519                    q8_range2_sdot(c.bytes, c.row_scale, &a.0, &a.1, c.cols, o1, o2, start, end)
2520                }
2521            });
2522            let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
2523                std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
2524            pool.run_many(&parts);
2525            return;
2526        }
2527        #[cfg(target_arch = "x86_64")]
2528        if avx2_a8w8_enabled() {
2529            let acts: [(SplitAct, SplitAct); N] =
2530                std::array::from_fn(|i| (split_act(&ctxs[i].xs1), split_act(&ctxs[i].xs2)));
2531            let closures: [_; N] = std::array::from_fn(|i| {
2532                let (c, a, o1, o2) = (&ctxs[i], &acts[i], p1[i], p2[i]);
2533                move |start: usize, end: usize| {
2534                    q8_range2_avx2(c.bytes, c.row_scale, &a.0, &a.1, c.cols, o1, o2, start, end)
2535                }
2536            });
2537            let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
2538                std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
2539            pool.run_many(&parts);
2540            return;
2541        }
2542        let closures: [_; N] = std::array::from_fn(|i| {
2543            let (c, o1, o2) = (&ctxs[i], p1[i], p2[i]);
2544            move |start: usize, end: usize| {
2545                q8_range2_f32(
2546                    c.bytes,
2547                    c.row_scale,
2548                    &c.xs1,
2549                    &c.xs2,
2550                    c.cols,
2551                    o1,
2552                    o2,
2553                    start,
2554                    end,
2555                )
2556            }
2557        });
2558        let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
2559            std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
2560        pool.run_many(&parts);
2561    }
2562
2563    /// Fused gate+up matvec with SiLU·mul: for each row r, computes
2564    /// `silu(gate·x) * (up·x)` and writes to `out[r]`. ONE pool dispatch,
2565    /// no intermediate g/u buffers, no separate silu pass. Falls back
2566    /// (returns false) for unsupported dtype combos.
2567    pub fn matvec_silu_mul(
2568        gate: &QTensor,
2569        up: &QTensor,
2570        x: &[f32],
2571        out: &mut [f32],
2572        pool: Option<&Pool>,
2573    ) -> bool {
2574        let inter = gate.rows();
2575        debug_assert_eq!(up.rows(), inter);
2576        debug_assert_eq!(out.len(), inter);
2577        debug_assert_eq!(gate.cols(), up.cols());
2578        if !a8w8_enabled() {
2579            return false;
2580        }
2581        let act = split_act(x);
2582        let act = &act;
2583        let x_ref = x;
2584        let out_addr = SendMut(out.as_mut_ptr());
2585
2586        match (gate, up) {
2587            // Q4Block gate + Q4Block up (most common mobile q4 models)
2588            (
2589                Self::Mapped {
2590                    dtype: TensorDtype::Q4Block,
2591                    ..
2592                },
2593                Self::Mapped {
2594                    dtype: TensorDtype::Q4Block,
2595                    ..
2596                },
2597            ) => {
2598                let (gp, gs) = q4_split(gate.quant_bytes(), gate.rows(), gate.cols());
2599                let (up_p, up_s) = q4_split(up.quant_bytes(), up.rows(), up.cols());
2600                let gpr = gate.cols() / GROUP_SIZE;
2601                let cols = gate.cols();
2602                let run = move |start: usize, end: usize| {
2603                    for r in start..end {
2604                        let mut gv = dot_q4_row_i8(gp, gs, r * gpr, gpr, &act.xq) * act.sx;
2605                        let mut uv = dot_q4_row_i8(up_p, up_s, r * gpr, gpr, &act.xq) * act.sx;
2606                        for &(j, xv) in &act.outliers {
2607                            let flat = r * cols + j;
2608                            let gb = gp[flat / 2];
2609                            let gn = if flat & 1 == 0 { gb & 0x0F } else { gb >> 4 };
2610                            let gsc = f16_to_f32(u16::from_le_bytes([
2611                                gs[(flat / GROUP_SIZE) * 2],
2612                                gs[(flat / GROUP_SIZE) * 2 + 1],
2613                            ]));
2614                            gv += ((gn as i32 - 8) as f32) * gsc * xv;
2615                            let ub = up_p[flat / 2];
2616                            let un = if flat & 1 == 0 { ub & 0x0F } else { ub >> 4 };
2617                            let usc = f16_to_f32(u16::from_le_bytes([
2618                                up_s[(flat / GROUP_SIZE) * 2],
2619                                up_s[(flat / GROUP_SIZE) * 2 + 1],
2620                            ]));
2621                            uv += ((un as i32 - 8) as f32) * usc * xv;
2622                        }
2623                        let silu_g = gv / (1.0 + (-gv).exp());
2624                        // SAFETY: disjoint row ranges per worker.
2625                        unsafe { *out_addr.at(r) = silu_g * uv };
2626                    }
2627                };
2628                dispatch_rows(pool, inter, &run);
2629                true
2630            }
2631            // Q4Tiled gate + Q4Tiled up — one row pass, both tile
2632            // streams sequential, silu·mul fused (same per-row math as
2633            // `q4t_matvec`).
2634            (
2635                Self::Mapped {
2636                    dtype: TensorDtype::Q4Tiled,
2637                    ..
2638                },
2639                Self::Mapped {
2640                    dtype: TensorDtype::Q4Tiled,
2641                    ..
2642                },
2643            ) => {
2644                let g_bytes = gate.quant_bytes();
2645                let u_bytes = up.quant_bytes();
2646                let gpr = gate.cols() / GROUP_SIZE;
2647                let run = move |start: usize, end: usize| {
2648                    for r in start..end {
2649                        let mut gv = dot_q4t_row_i8(g_bytes, r, gpr, &act.xq) * act.sx;
2650                        let mut uv = dot_q4t_row_i8(u_bytes, r, gpr, &act.xq) * act.sx;
2651                        for &(j, xv) in &act.outliers {
2652                            let (w, s) = q4t_outlier(g_bytes, r, gpr, j);
2653                            gv += w * s * xv;
2654                            let (w, s) = q4t_outlier(u_bytes, r, gpr, j);
2655                            uv += w * s * xv;
2656                        }
2657                        let silu_g = gv / (1.0 + (-gv).exp());
2658                        // SAFETY: disjoint row ranges per worker.
2659                        unsafe { *out_addr.at(r) = silu_g * uv };
2660                    }
2661                };
2662                dispatch_rows(pool, inter, &run);
2663                true
2664            }
2665            // Q4TiledP gate + Q4TiledP up — the same fused row pass, with
2666            // each row's two ladders built once and spent on both streams.
2667            (
2668                Self::Mapped {
2669                    dtype: TensorDtype::Q4TiledP,
2670                    ..
2671                },
2672                Self::Mapped {
2673                    dtype: TensorDtype::Q4TiledP,
2674                    ..
2675                },
2676            ) => {
2677                let cols = gate.cols();
2678                let gpr = cols / GROUP_SIZE;
2679                let gv_view = Q4tpView::new(gate.quant_bytes(), inter, cols);
2680                let uv_view = Q4tpView::new(up.quant_bytes(), inter, cols);
2681                let run = |start: usize, end: usize| {
2682                    let (mut gsc, mut usc) = (vec![0f32; gpr], vec![0f32; gpr]);
2683                    for r in start..end {
2684                        gv_view.scales_into(r, gpr, &mut gsc);
2685                        uv_view.scales_into(r, gpr, &mut usc);
2686                        let mut gv = dot_q4tp_row_i8(gv_view.nib, r, gpr, &act.xq, &gsc) * act.sx;
2687                        let mut uv = dot_q4tp_row_i8(uv_view.nib, r, gpr, &act.xq, &usc) * act.sx;
2688                        for &(j, xv) in &act.outliers {
2689                            let (w, s) = q4tp_outlier(gv_view.nib, r, gpr, j, &gsc);
2690                            gv += w * s * xv;
2691                            let (w, s) = q4tp_outlier(uv_view.nib, r, gpr, j, &usc);
2692                            uv += w * s * xv;
2693                        }
2694                        let silu_g = gv / (1.0 + (-gv).exp());
2695                        // SAFETY: disjoint row ranges per worker.
2696                        unsafe { *out_addr.at(r) = silu_g * uv };
2697                    }
2698                };
2699                dispatch_rows(pool, inter, &run);
2700                true
2701            }
2702            // Q1 gate + Q1 up — one row pass over both sign streams,
2703            // silu·mul fused (the per-row math of `q1_range_a8w8`); the
2704            // activation group sums are shared by both streams. Without
2705            // this arm a q1 dense FFN paid two dispatches + a combine
2706            // loop — the exact barrier this function exists to remove.
2707            (
2708                Self::Mapped {
2709                    dtype: TensorDtype::Q1,
2710                    ..
2711                },
2712                Self::Mapped {
2713                    dtype: TensorDtype::Q1,
2714                    ..
2715                },
2716            ) => {
2717                let g_bytes = gate.quant_bytes();
2718                let u_bytes = up.quant_bytes();
2719                let gpr = gate.cols() / GROUP_SIZE;
2720                let gsum = q1_group_sums(&act.xq, gpr);
2721                let gsum = &gsum;
2722                let run = move |start: usize, end: usize| {
2723                    for r in start..end {
2724                        let mut gv = dot_q1_row_i8(g_bytes, r, gpr, &act.xq, gsum) * act.sx;
2725                        let mut uv = dot_q1_row_i8(u_bytes, r, gpr, &act.xq, gsum) * act.sx;
2726                        for &(j, xv) in &act.outliers {
2727                            let (w, s) = q1_outlier(g_bytes, r, gpr, j);
2728                            gv += w * s * xv;
2729                            let (w, s) = q1_outlier(u_bytes, r, gpr, j);
2730                            uv += w * s * xv;
2731                        }
2732                        let silu_g = gv / (1.0 + (-gv).exp());
2733                        // SAFETY: disjoint row ranges per worker.
2734                        unsafe { *out_addr.at(r) = silu_g * uv };
2735                    }
2736                };
2737                dispatch_rows(pool, inter, &run);
2738                true
2739            }
2740            // Q2TiledP gate + Q2TiledP up — the 2-bit expert pair (MoE
2741            // FFNs of the W2 class): one row pass, both ladders built
2742            // once, integer code dots with shared group sums.
2743            (
2744                Self::Mapped {
2745                    dtype: TensorDtype::Q2TiledP,
2746                    ..
2747                },
2748                Self::Mapped {
2749                    dtype: TensorDtype::Q2TiledP,
2750                    ..
2751                },
2752            ) => {
2753                let cols = gate.cols();
2754                let gpr = cols / GROUP_SIZE;
2755                let gv_view = Q4tpView::new_q2(gate.quant_bytes(), inter, cols);
2756                let uv_view = Q4tpView::new_q2(up.quant_bytes(), inter, cols);
2757                let gsum = q1_group_sums(&act.xq, gpr);
2758                let gsum = &gsum;
2759                let run = move |start: usize, end: usize| {
2760                    let (mut gsc, mut usc) = (vec![0f32; gpr], vec![0f32; gpr]);
2761                    for r in start..end {
2762                        gv_view.scales_into(r, gpr, &mut gsc);
2763                        uv_view.scales_into(r, gpr, &mut usc);
2764                        let mut gv =
2765                            dot_q2tp_row_i8(gv_view.nib, r, gpr, &act.xq, gsum, &gsc) * act.sx;
2766                        let mut uv =
2767                            dot_q2tp_row_i8(uv_view.nib, r, gpr, &act.xq, gsum, &usc) * act.sx;
2768                        for &(j, xv) in &act.outliers {
2769                            let (w, s) = q2tp_outlier(gv_view.nib, r, gpr, j, &gsc);
2770                            gv += w * s * xv;
2771                            let (w, s) = q2tp_outlier(uv_view.nib, r, gpr, j, &usc);
2772                            uv += w * s * xv;
2773                        }
2774                        let silu_g = gv / (1.0 + (-gv).exp());
2775                        // SAFETY: disjoint row ranges per worker.
2776                        unsafe { *out_addr.at(r) = silu_g * uv };
2777                    }
2778                };
2779                dispatch_rows(pool, inter, &run);
2780                true
2781            }
2782            // Q8Row gate + Q8Row up — one row pass over both i8 streams.
2783            // Q8_2f stays out on purpose: its column field prescales the
2784            // activations PER TENSOR, which breaks this fn's shared
2785            // split_act contract — it keeps the two-dispatch path.
2786            (
2787                Self::Mapped {
2788                    dtype: TensorDtype::Q8Row,
2789                    row_scale: g_rs,
2790                    ..
2791                },
2792                Self::Mapped {
2793                    dtype: TensorDtype::Q8Row,
2794                    row_scale: u_rs,
2795                    ..
2796                },
2797            ) => {
2798                let g_bytes = gate.quant_bytes();
2799                let u_bytes = up.quant_bytes();
2800                let cols = gate.cols();
2801                let run = move |start: usize, end: usize| {
2802                    for r in start..end {
2803                        let gv = q8_row_dot(&g_bytes[r * cols..(r + 1) * cols], act) * g_rs[r];
2804                        let uv = q8_row_dot(&u_bytes[r * cols..(r + 1) * cols], act) * u_rs[r];
2805                        let silu_g = gv / (1.0 + (-gv).exp());
2806                        // SAFETY: disjoint row ranges per worker.
2807                        unsafe { *out_addr.at(r) = silu_g * uv };
2808                    }
2809                };
2810                dispatch_rows(pool, inter, &run);
2811                true
2812            }
2813            // Q1T gate + Q1T up
2814            (
2815                Self::Mapped {
2816                    dtype: TensorDtype::Q1T,
2817                    ..
2818                },
2819                Self::Mapped {
2820                    dtype: TensorDtype::Q1T,
2821                    ..
2822                },
2823            ) => {
2824                const TILE: usize = cortiq_core::quant::Q1T_TILE;
2825                let g_bytes = gate.quant_bytes();
2826                let u_bytes = up.quant_bytes();
2827                let gpr = gate.cols() / GROUP_SIZE;
2828                let (g_rp, g_ent, g_ov) = q1t_overlay(g_bytes, inter * gpr * TILE, inter);
2829                let (u_rp, u_ent, u_ov) = q1t_overlay(u_bytes, inter * gpr * TILE, inter);
2830                let run = move |start: usize, end: usize| {
2831                    for r in start..end {
2832                        let mut gv = q1t_dot_row_i8(g_bytes, r, gpr, &act.xq) * act.sx;
2833                        let mut uv = q1t_dot_row_i8(u_bytes, r, gpr, &act.xq) * act.sx;
2834                        for &(j, xv) in &act.outliers {
2835                            gv += q1t_base_weight(g_bytes, r, gpr, j) * xv;
2836                            uv += q1t_base_weight(u_bytes, r, gpr, j) * xv;
2837                        }
2838                        gv += q1t_row_outlier_correction(g_bytes, r, g_rp, g_ent, g_ov, x_ref);
2839                        uv += q1t_row_outlier_correction(u_bytes, r, u_rp, u_ent, u_ov, x_ref);
2840                        let silu_g = gv / (1.0 + (-gv).exp());
2841                        // SAFETY: disjoint row ranges per worker.
2842                        unsafe { *out_addr.at(r) = silu_g * uv };
2843                    }
2844                };
2845                dispatch_rows(pool, inter, &run);
2846                true
2847            }
2848            _ => false,
2849        }
2850    }
2851
2852    /// Every routed expert's fused gate/up/SiLU under ONE pool dispatch.
2853    ///
2854    /// The per-expert path pays a pool barrier per expert per stage: at 9
2855    /// experts over 40 layers that is ~720 barriers a token, and a decode
2856    /// profile of Qwen3.6-35B-A3B showed the pool parked in
2857    /// `psynch_cvwait` about twice as long as it spent computing. Laying
2858    /// every expert's rows end-to-end in one virtual row space collapses
2859    /// the stage to a single dispatch. The per-row body is the
2860    /// single-expert q4tp arm verbatim, so outputs are bit-identical.
2861    ///
2862    /// `false` = something is outside the fused q4tp kernel (dtype, shape,
2863    /// or the `CMF_SDOT=0` exact contract); the caller walks the ordinary
2864    /// per-expert path.
2865    pub fn moe_gate_up_many(
2866        pairs: &[(&QTensor, &QTensor)],
2867        x: &[f32],
2868        outs: &mut [Vec<f32>],
2869        pool: Option<&Pool>,
2870    ) -> bool {
2871        if pairs.is_empty() || pairs.len() != outs.len() || !a8w8_enabled() {
2872            return false;
2873        }
2874        let inter = pairs[0].0.rows();
2875        let cols = pairs[0].0.cols();
2876        if cols % GROUP_SIZE != 0 {
2877            return false;
2878        }
2879        let gpr = cols / GROUP_SIZE;
2880        // Uniform layout across every routed pair: q4tp, or the 2-bit
2881        // profile's q2tp gate/up (the W2 class). Mixed sets refuse.
2882        let q2 = matches!(
2883            pairs[0].0,
2884            Self::Mapped {
2885                dtype: TensorDtype::Q2TiledP,
2886                ..
2887            }
2888        );
2889        let want = if q2 {
2890            TensorDtype::Q2TiledP
2891        } else {
2892            TensorDtype::Q4TiledP
2893        };
2894        let mut views = Vec::with_capacity(pairs.len() * 2);
2895        for ((g, u), o) in pairs.iter().zip(outs.iter()) {
2896            let both = matches!(g, Self::Mapped { dtype, .. } if *dtype == want)
2897                && matches!(u, Self::Mapped { dtype, .. } if *dtype == want);
2898            if !both
2899                || g.rows() != inter
2900                || u.rows() != inter
2901                || g.cols() != cols
2902                || u.cols() != cols
2903                || o.len() != inter
2904            {
2905                return false;
2906            }
2907            let mk = if q2 { Q4tpView::new_q2 } else { Q4tpView::new };
2908            views.push(mk(g.quant_bytes(), inter, cols));
2909            views.push(mk(u.quant_bytes(), inter, cols));
2910        }
2911        let act = split_act(x);
2912        let gsum = if q2 {
2913            q1_group_sums(&act.xq, gpr)
2914        } else {
2915            Vec::new()
2916        };
2917        let (act, gsum) = (&act, &gsum);
2918        let ptrs: Vec<SendMut> = outs.iter_mut().map(|o| SendMut(o.as_mut_ptr())).collect();
2919        let (views, ptrs) = (&views, &ptrs);
2920        let run = |start: usize, end: usize| {
2921            let (mut gsc, mut usc) = (vec![0f32; gpr], vec![0f32; gpr]);
2922            for flat in start..end {
2923                let (e, r) = (flat / inter, flat % inter);
2924                let gv_view = &views[e * 2];
2925                let uv_view = &views[e * 2 + 1];
2926                gv_view.scales_into(r, gpr, &mut gsc);
2927                uv_view.scales_into(r, gpr, &mut usc);
2928                let (mut gv, mut uv) = if q2 {
2929                    (
2930                        dot_q2tp_row_i8(gv_view.nib, r, gpr, &act.xq, gsum, &gsc) * act.sx,
2931                        dot_q2tp_row_i8(uv_view.nib, r, gpr, &act.xq, gsum, &usc) * act.sx,
2932                    )
2933                } else {
2934                    (
2935                        dot_q4tp_row_i8(gv_view.nib, r, gpr, &act.xq, &gsc) * act.sx,
2936                        dot_q4tp_row_i8(uv_view.nib, r, gpr, &act.xq, &usc) * act.sx,
2937                    )
2938                };
2939                for &(j, xv) in &act.outliers {
2940                    let (og, ou) = if q2 {
2941                        (
2942                            q2tp_outlier(gv_view.nib, r, gpr, j, &gsc),
2943                            q2tp_outlier(uv_view.nib, r, gpr, j, &usc),
2944                        )
2945                    } else {
2946                        (
2947                            q4tp_outlier(gv_view.nib, r, gpr, j, &gsc),
2948                            q4tp_outlier(uv_view.nib, r, gpr, j, &usc),
2949                        )
2950                    };
2951                    gv += og.0 * og.1 * xv;
2952                    uv += ou.0 * ou.1 * xv;
2953                }
2954                let silu_g = gv / (1.0 + (-gv).exp());
2955                // SAFETY: one worker owns each (expert, row) pair.
2956                unsafe { *ptrs[e].at(r) = silu_g * uv };
2957            }
2958        };
2959        dispatch_rows(pool, pairs.len() * inter, &run);
2960        true
2961    }
2962
2963    /// Every routed expert's down projection, weighted and summed into
2964    /// `out`, under ONE pool dispatch.
2965    ///
2966    /// Partitioned by OUTPUT row rather than by expert: each row is owned
2967    /// by a single worker, so the experts are summed in the caller's order
2968    /// — the same sequence of f32 adds the serial `out[i] += w·eo[i]` loop
2969    /// performs, hence bit-identical. Partitioning by expert instead would
2970    /// race on the shared accumulator.
2971    pub fn moe_down_many(
2972        downs: &[&QTensor],
2973        gs: &[Vec<f32>],
2974        weights: &[f32],
2975        out: &mut [f32],
2976        pool: Option<&Pool>,
2977    ) -> bool {
2978        if downs.is_empty()
2979            || downs.len() != gs.len()
2980            || downs.len() != weights.len()
2981            || !a8w8_enabled()
2982        {
2983            return false;
2984        }
2985        let rows = out.len();
2986        let cols = downs[0].cols();
2987        if cols % GROUP_SIZE != 0 {
2988            return false;
2989        }
2990        let gpr = cols / GROUP_SIZE;
2991        let mut views = Vec::with_capacity(downs.len());
2992        for (d, g) in downs.iter().zip(gs.iter()) {
2993            if !matches!(
2994                d,
2995                Self::Mapped {
2996                    dtype: TensorDtype::Q4TiledP,
2997                    ..
2998                }
2999            ) || d.rows() != rows
3000                || d.cols() != cols
3001                || g.len() != cols
3002            {
3003                return false;
3004            }
3005            views.push(Q4tpView::new(d.quant_bytes(), rows, cols));
3006        }
3007        // One int8 split per expert — the activation vectors differ.
3008        let acts: Vec<SplitAct> = gs.iter().map(|g| split_act(g)).collect();
3009        // Partitioned by OUTPUT row, with the experts folded inside: each
3010        // row is owned by one worker, so they are summed in the caller's
3011        // order — the same f32 sequence the serial `out[i] += w·eo[i]`
3012        // loop produces. Partitioning by expert instead would either race
3013        // on the accumulator or need a scratch plane and a second pass;
3014        // measured, that variant was a wash, so this keeps the simpler
3015        // shape.
3016        let out_addr = SendMut(out.as_mut_ptr());
3017        let (views, acts, weights) = (&views, &acts, &weights);
3018        let run = |start: usize, end: usize| {
3019            let mut sc = vec![0f32; gpr];
3020            for r in start..end {
3021                let mut acc = 0f32;
3022                for (e, v) in views.iter().enumerate() {
3023                    v.scales_into(r, gpr, &mut sc);
3024                    let a = &acts[e];
3025                    let mut d = dot_q4tp_row_i8(v.nib, r, gpr, &a.xq, &sc) * a.sx;
3026                    for &(j, xv) in &a.outliers {
3027                        let (w, s) = q4tp_outlier(v.nib, r, gpr, j, &sc);
3028                        d += w * s * xv;
3029                    }
3030                    acc += weights[e] * d;
3031                }
3032                // SAFETY: disjoint row ranges per worker.
3033                unsafe { *out_addr.at(r) = acc };
3034            }
3035        };
3036        dispatch_rows(pool, rows, &run);
3037        true
3038    }
3039}
3040
3041/// Batched q8 kernel: same math as qmatvec, the row makes a single
3042/// pass from memory for the whole batch.
3043/// Accelerate CBLAS — the Apple AMX matrix units, the same engine
3044/// llama.cpp's `-ngl 0` prefill rides via ggml-blas.
3045#[cfg(target_os = "macos")]
3046mod accel_blas {
3047    #[link(name = "Accelerate", kind = "framework")]
3048    unsafe extern "C" {
3049        pub fn cblas_sgemm(
3050            order: i32,
3051            trans_a: i32,
3052            trans_b: i32,
3053            m: i32,
3054            n: i32,
3055            k: i32,
3056            alpha: f32,
3057            a: *const f32,
3058            lda: i32,
3059            b: *const f32,
3060            ldb: i32,
3061            beta: f32,
3062            c: *mut f32,
3063            ldc: i32,
3064        );
3065    }
3066}
3067
3068#[cfg(target_os = "macos")]
3069pub(crate) fn accel_gemm_enabled() -> bool {
3070    static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
3071    *ON.get_or_init(|| std::env::var("CMF_ACCEL").map(|v| v != "0").unwrap_or(true))
3072}
3073
3074/// Off macOS the "accel" GEMM is the portable NEON micro-kernel below —
3075/// same entry point, so the batched-attention path opens on mobile.
3076#[cfg(all(target_arch = "aarch64", not(target_os = "macos")))]
3077pub(crate) fn accel_gemm_enabled() -> bool {
3078    true
3079}
3080
3081/// Portable NEON f32 GEMM (row-major, optional Bᵀ): a 4×8 fmla
3082/// micro-kernel with A broadcast against B panels — the mobile stand-in
3083/// for Accelerate in the batched causal attention (QKᵀ and P·V). Not a
3084/// BLAS: shapes here are the attention panels (m ≤ heads·chunk,
3085/// k = head_dim or context), and the goal is removing the per-position
3086/// quadratic wall, not peak GEMM.
3087#[cfg(target_arch = "aarch64")]
3088#[allow(clippy::too_many_arguments)]
3089pub(crate) fn neon_gemm_rm(
3090    m: usize,
3091    n: usize,
3092    k: usize,
3093    alpha: f32,
3094    a: &[f32],
3095    lda: usize,
3096    b_mat: &[f32],
3097    ldb: usize,
3098    b_rows_are_n: bool,
3099    c: &mut [f32],
3100    ldc: usize,
3101) {
3102    debug_assert!(a.len() >= (m - 1) * lda + k);
3103    debug_assert!(c.len() >= (m - 1) * ldc + n);
3104    // SAFETY: bounds asserted above; NEON is baseline on aarch64.
3105    unsafe {
3106        use core::arch::aarch64::*;
3107        let mut i = 0usize;
3108        while i < m {
3109            let mi = (m - i).min(4);
3110            let mut j = 0usize;
3111            while j < n {
3112                let nj = (n - j).min(8);
3113                if mi == 4 && nj == 8 {
3114                    let (mut c0a, mut c0b) = (vdupq_n_f32(0.0), vdupq_n_f32(0.0));
3115                    let (mut c1a, mut c1b) = (vdupq_n_f32(0.0), vdupq_n_f32(0.0));
3116                    let (mut c2a, mut c2b) = (vdupq_n_f32(0.0), vdupq_n_f32(0.0));
3117                    let (mut c3a, mut c3b) = (vdupq_n_f32(0.0), vdupq_n_f32(0.0));
3118                    for p in 0..k {
3119                        let (b0, b1) = if b_rows_are_n {
3120                            // B is [n, k]: column p of Bᵀ = element p of
3121                            // eight consecutive B rows — gathered.
3122                            let base = b_mat.as_ptr().add(j * ldb + p);
3123                            let g = |o: usize| *base.add(o * ldb);
3124                            ([g(0), g(1), g(2), g(3)], [g(4), g(5), g(6), g(7)])
3125                        } else {
3126                            let base = b_mat.as_ptr().add(p * ldb + j);
3127                            (
3128                                [*base, *base.add(1), *base.add(2), *base.add(3)],
3129                                [*base.add(4), *base.add(5), *base.add(6), *base.add(7)],
3130                            )
3131                        };
3132                        let bv0 = vld1q_f32(b0.as_ptr());
3133                        let bv1 = vld1q_f32(b1.as_ptr());
3134                        let a0 = vdupq_n_f32(*a.as_ptr().add(i * lda + p));
3135                        let a1 = vdupq_n_f32(*a.as_ptr().add((i + 1) * lda + p));
3136                        let a2 = vdupq_n_f32(*a.as_ptr().add((i + 2) * lda + p));
3137                        let a3 = vdupq_n_f32(*a.as_ptr().add((i + 3) * lda + p));
3138                        c0a = vfmaq_f32(c0a, a0, bv0);
3139                        c0b = vfmaq_f32(c0b, a0, bv1);
3140                        c1a = vfmaq_f32(c1a, a1, bv0);
3141                        c1b = vfmaq_f32(c1b, a1, bv1);
3142                        c2a = vfmaq_f32(c2a, a2, bv0);
3143                        c2b = vfmaq_f32(c2b, a2, bv1);
3144                        c3a = vfmaq_f32(c3a, a3, bv0);
3145                        c3b = vfmaq_f32(c3b, a3, bv1);
3146                    }
3147                    let al = vdupq_n_f32(alpha);
3148                    for (r, (ca, cb)) in [(c0a, c0b), (c1a, c1b), (c2a, c2b), (c3a, c3b)]
3149                        .iter()
3150                        .enumerate()
3151                    {
3152                        let dst = c.as_mut_ptr().add((i + r) * ldc + j);
3153                        vst1q_f32(dst, vmulq_f32(*ca, al));
3154                        vst1q_f32(dst.add(4), vmulq_f32(*cb, al));
3155                    }
3156                } else {
3157                    for r in 0..mi {
3158                        for q in 0..nj {
3159                            let mut acc = 0f32;
3160                            for p in 0..k {
3161                                let bv = if b_rows_are_n {
3162                                    b_mat[(j + q) * ldb + p]
3163                                } else {
3164                                    b_mat[p * ldb + j + q]
3165                                };
3166                                acc += a[(i + r) * lda + p] * bv;
3167                            }
3168                            c[(i + r) * ldc + j + q] = acc * alpha;
3169                        }
3170                    }
3171                }
3172                j += nj;
3173            }
3174            i += mi;
3175        }
3176    }
3177}
3178
3179/// Off-macOS aarch64: the batched attention rides the NEON micro-GEMM.
3180#[cfg(all(target_arch = "aarch64", not(target_os = "macos")))]
3181#[allow(clippy::too_many_arguments)]
3182pub(crate) fn sgemm_rm(
3183    m: usize,
3184    n: usize,
3185    k: usize,
3186    alpha: f32,
3187    a: &[f32],
3188    lda: usize,
3189    b_mat: &[f32],
3190    ldb: usize,
3191    b_rows_are_n: bool,
3192    c: &mut [f32],
3193    ldc: usize,
3194) {
3195    neon_gemm_rm(m, n, k, alpha, a, lda, b_mat, ldb, b_rows_are_n, c, ldc);
3196}
3197
3198/// Row-major f32 GEMM, exposed for offline tools (the AWNP pass builds a
3199/// per-layer projection and applies it to every expert; a naive triple loop
3200/// would turn a two-minute job into half an hour).
3201#[allow(clippy::too_many_arguments)]
3202pub fn sgemm_public(
3203    m: usize,
3204    n: usize,
3205    k: usize,
3206    alpha: f32,
3207    a: &[f32],
3208    lda: usize,
3209    b_mat: &[f32],
3210    ldb: usize,
3211    b_rows_are_n: bool,
3212    c: &mut [f32],
3213    ldc: usize,
3214) {
3215    #[cfg(any(target_os = "macos", target_arch = "aarch64"))]
3216    {
3217        sgemm_rm(m, n, k, alpha, a, lda, b_mat, ldb, b_rows_are_n, c, ldc);
3218    }
3219    // x86 without Accelerate has no sgemm_rm: the specialized paths there are
3220    // quantized kernels, not an f32 GEMM. Only the offline AWNP pass reaches
3221    // this, so correctness matters and throughput does not — a triple loop is
3222    // the honest fallback rather than a reason to make the tool macOS-only.
3223    #[cfg(not(any(target_os = "macos", target_arch = "aarch64")))]
3224    {
3225        for i in 0..m {
3226            for j in 0..n {
3227                let mut acc = 0f32;
3228                for p in 0..k {
3229                    let bv = if b_rows_are_n {
3230                        b_mat[j * ldb + p]
3231                    } else {
3232                        b_mat[p * ldb + j]
3233                    };
3234                    acc += a[i * lda + p] * bv;
3235                }
3236                c[i * ldc + j] = alpha * acc;
3237            }
3238        }
3239    }
3240}
3241
3242/// Row-major f32 GEMM on Accelerate: C[m,n] = alpha·A[m,k] × B(ᵀ).
3243/// `b_rows_are_n` = true multiplies by Bᵀ where B is stored [n, k].
3244#[cfg(target_os = "macos")]
3245#[allow(clippy::too_many_arguments)]
3246pub(crate) fn sgemm_rm(
3247    m: usize,
3248    n: usize,
3249    k: usize,
3250    alpha: f32,
3251    a: &[f32],
3252    lda: usize,
3253    b_mat: &[f32],
3254    ldb: usize,
3255    b_rows_are_n: bool,
3256    c: &mut [f32],
3257    ldc: usize,
3258) {
3259    debug_assert!(a.len() >= (m - 1) * lda + k);
3260    debug_assert!(c.len() >= (m - 1) * ldc + n);
3261    // Test hook: route the attention GEMMs through the portable NEON
3262    // micro-kernel ON APPLE SILICON — how the mobile batched attend is
3263    // measured without a phone in the loop. (Intel macOS has no NEON —
3264    // the hook is a no-op there, Accelerate continues below.)
3265    #[cfg(target_arch = "aarch64")]
3266    if std::env::var("CMF_FORCE_NEON_GEMM")
3267        .map(|v| v == "1")
3268        .unwrap_or(false)
3269    {
3270        return neon_gemm_rm(m, n, k, alpha, a, lda, b_mat, ldb, b_rows_are_n, c, ldc);
3271    }
3272    unsafe {
3273        accel_blas::cblas_sgemm(
3274            101, // RowMajor
3275            111, // NoTrans A
3276            if b_rows_are_n { 112 } else { 111 },
3277            m as i32,
3278            n as i32,
3279            k as i32,
3280            alpha,
3281            a.as_ptr(),
3282            lda as i32,
3283            b_mat.as_ptr(),
3284            ldb as i32,
3285            0.0,
3286            c.as_mut_ptr(),
3287            ldc as i32,
3288        );
3289    }
3290}
3291
3292/// Prefill GEMM through Accelerate (macOS): dequantize q8 rows into
3293/// f32 tiles (scale folded in, pool-parallel) and multiply each tile
3294/// on the AMX with one row-major sgemm. Tiles live in cache, weights
3295/// stream once. Numerics are f32-GEMM (not the int8 dot): prefill
3296/// logits shift within f32 rounding — tolerance-class, like every
3297/// reduction-order change; decode (M=1) never takes this path.
3298#[cfg(target_os = "macos")]
3299fn qmatmat_accel(
3300    q: &[u8],
3301    row_scale: &[f32],
3302    pre: &[std::borrow::Cow<'_, [f32]>],
3303    rows: usize,
3304    cols: usize,
3305    out: &mut [f32],
3306    pool: Option<&Pool>,
3307) {
3308    // NOTE: double-buffering the dequant against the sgemm (a scoped
3309    // thread driving the pool on tile k+1 while the caller multiplies
3310    // tile k) was tried and LOST ~6%: Accelerate's sgemm is itself
3311    // multithreaded, and the dequant workers just steal its cores.
3312    const TR: usize = 2048;
3313    let b = pre.len();
3314    thread_local! {
3315        static XPANEL: std::cell::RefCell<Vec<f32>> = const { std::cell::RefCell::new(Vec::new()) };
3316        static WTILE: std::cell::RefCell<Vec<f32>> = const { std::cell::RefCell::new(Vec::new()) };
3317    }
3318    XPANEL.with(|xp| {
3319        WTILE.with(|wt| {
3320            let mut xpanel = xp.borrow_mut();
3321            xpanel.clear();
3322            for x in pre {
3323                xpanel.extend_from_slice(x);
3324            }
3325            let mut wtile = wt.borrow_mut();
3326            wtile.resize(TR * cols, 0.0);
3327            let mut r0 = 0usize;
3328            while r0 < rows {
3329                let tr = TR.min(rows - r0);
3330                // Dequant the tile (scale folded) — pool-parallel.
3331                let wt_addr = SendMut(wtile.as_mut_ptr());
3332                let run = |start: usize, end: usize| {
3333                    for r in start..end {
3334                        let row = &q[(r0 + r) * cols..(r0 + r + 1) * cols];
3335                        let s = row_scale[r0 + r];
3336                        // SAFETY: workers cover disjoint r ranges.
3337                        let dst =
3338                            unsafe { std::slice::from_raw_parts_mut(wt_addr.at(r * cols), cols) };
3339                        for (d, &v) in dst.iter_mut().zip(row) {
3340                            *d = (v as i8) as f32 * s;
3341                        }
3342                    }
3343                };
3344                dispatch_rows(pool, tr, &run);
3345                // C[b, tr] (at column r0 of out[b, rows]) = X · Wtileᵀ
3346                unsafe {
3347                    accel_blas::cblas_sgemm(
3348                        101, // RowMajor
3349                        111, // NoTrans A
3350                        112, // Trans B
3351                        b as i32,
3352                        tr as i32,
3353                        cols as i32,
3354                        1.0,
3355                        xpanel.as_ptr(),
3356                        cols as i32,
3357                        wtile.as_ptr(),
3358                        cols as i32,
3359                        0.0,
3360                        out.as_mut_ptr().add(r0),
3361                        rows as i32,
3362                    );
3363                }
3364                r0 += tr;
3365            }
3366        })
3367    });
3368}
3369
3370fn qmatmat(
3371    q: &[u8],
3372    row_scale: &[f32],
3373    pre: &[std::borrow::Cow<'_, [f32]>],
3374    rows: usize,
3375    cols: usize,
3376    out: &mut [f32],
3377    pool: Option<&Pool>,
3378) {
3379    let b = pre.len();
3380    debug_assert_eq!(out.len(), b * rows);
3381    // Big prefill batches ride the AMX (roadmap PR3): the row×batch
3382    // SDOT loop below peaks near the CPU's dot throughput, an order
3383    // below the matrix units. Small tensors and tiny test models stay
3384    // on the exact integer path.
3385    #[cfg(target_os = "macos")]
3386    if b >= 8 && rows * cols >= 500_000 && accel_gemm_enabled() {
3387        qmatmat_accel(q, row_scale, pre, rows, cols, out, pool);
3388        return;
3389    }
3390    #[cfg(target_arch = "aarch64")]
3391    if sdot_enabled() {
3392        let acts: Vec<SplitAct> = pre.iter().map(|x| split_act(x)).collect();
3393        let out_addr = SendMut(out.as_mut_ptr());
3394        // Blocked 2×4 (mobile prefill: no AMX to fall back on — this
3395        // path IS the ARM prefill GEMM off Apple silicon).
3396        let blocked_ok = blocked_enabled();
3397        let use_i8mm = i8mm_enabled();
3398        if blocked_ok {
3399            let run = |start: usize, end: usize| {
3400                let mut o = start;
3401                while o < end {
3402                    if o + 2 <= end {
3403                        let r0 = &q[o * cols..(o + 1) * cols];
3404                        let r1 = &q[(o + 1) * cols..(o + 2) * cols];
3405                        let mut bi = 0usize;
3406                        while bi + 4 <= acts.len() {
3407                            let xs = [
3408                                acts[bi].xq.as_slice(),
3409                                acts[bi + 1].xq.as_slice(),
3410                                acts[bi + 2].xq.as_slice(),
3411                                acts[bi + 3].xq.as_slice(),
3412                            ];
3413                            let d = if use_i8mm {
3414                                unsafe { dot_i8_smmla_2x4(r0, r1, xs) }
3415                            } else {
3416                                unsafe { dot_i8_sdot_2x4(r0, r1, xs) }
3417                            };
3418                            for (r, row) in [r0, r1].into_iter().enumerate() {
3419                                for k in 0..4 {
3420                                    let act = &acts[bi + k];
3421                                    let mut v = d[r][k] as f32 * act.sx;
3422                                    for &(j, xv) in &act.outliers {
3423                                        v += (row[j] as i8) as f32 * xv;
3424                                    }
3425                                    unsafe {
3426                                        *out_addr.at((bi + k) * rows + o + r) = v * row_scale[o + r]
3427                                    };
3428                                }
3429                            }
3430                            bi += 4;
3431                        }
3432                        while bi < acts.len() {
3433                            for (r, row) in [r0, r1].into_iter().enumerate() {
3434                                let v = row_dot_sdot(row, &acts[bi]) * row_scale[o + r];
3435                                unsafe { *out_addr.at(bi * rows + o + r) = v };
3436                            }
3437                            bi += 1;
3438                        }
3439                        o += 2;
3440                    } else {
3441                        let row = &q[o * cols..(o + 1) * cols];
3442                        for (bi, act) in acts.iter().enumerate() {
3443                            let v = row_dot_sdot(row, act) * row_scale[o];
3444                            unsafe { *out_addr.at(bi * rows + o) = v };
3445                        }
3446                        o += 1;
3447                    }
3448                }
3449            };
3450            dispatch_rows(pool, rows, &run);
3451            return;
3452        }
3453        let run = |start: usize, end: usize| {
3454            for o in start..end {
3455                let row = &q[o * cols..(o + 1) * cols];
3456                for (bi, act) in acts.iter().enumerate() {
3457                    let v = row_dot_sdot(row, act) * row_scale[o];
3458                    unsafe { *out_addr.at(bi * rows + o) = v };
3459                }
3460            }
3461        };
3462        dispatch_rows(pool, rows, &run);
3463        return;
3464    }
3465    // x86 A8W8 batch. Non-VNNI parts take the BLOCKED 2×4 kernel
3466    // (roadmap P0: two weight rows' abs() stay in registers across four
3467    // activation streams); VNNI machines keep the per-row bias-trick
3468    // dot, which is already throughput-bound there.
3469    #[cfg(target_arch = "x86_64")]
3470    if avx2_a8w8_enabled() {
3471        let acts: Vec<SplitAct> = pre.iter().map(|x| split_act(x)).collect();
3472        let out_addr = SendMut(out.as_mut_ptr());
3473        // CMF_X86_BLOCKED=0 forces the per-row path (paired in-process
3474        // A/B on noisy shared-vCPU hosts).
3475        let blocked_ok = blocked_enabled();
3476        if !avx512vnni_enabled() && blocked_ok {
3477            let run = |start: usize, end: usize| {
3478                let mut o = start;
3479                while o < end {
3480                    if o + 2 <= end {
3481                        let r0 = &q[o * cols..(o + 1) * cols];
3482                        let r1 = &q[(o + 1) * cols..(o + 2) * cols];
3483                        let mut bi = 0usize;
3484                        while bi + 4 <= acts.len() {
3485                            let xs = [
3486                                acts[bi].xq.as_slice(),
3487                                acts[bi + 1].xq.as_slice(),
3488                                acts[bi + 2].xq.as_slice(),
3489                                acts[bi + 3].xq.as_slice(),
3490                            ];
3491                            let d = unsafe { dot_i8_i8_avx2_2x4(r0, r1, xs) };
3492                            for (r, row) in [r0, r1].into_iter().enumerate() {
3493                                for k in 0..4 {
3494                                    let act = &acts[bi + k];
3495                                    let mut v = d[r][k] as f32 * act.sx;
3496                                    for &(j, xv) in &act.outliers {
3497                                        v += (row[j] as i8) as f32 * xv;
3498                                    }
3499                                    unsafe {
3500                                        *out_addr.at((bi + k) * rows + o + r) = v * row_scale[o + r]
3501                                    };
3502                                }
3503                            }
3504                            bi += 4;
3505                        }
3506                        while bi < acts.len() {
3507                            for (r, row) in [r0, r1].into_iter().enumerate() {
3508                                let v = row_dot_avx2(row, &acts[bi]) * row_scale[o + r];
3509                                unsafe { *out_addr.at(bi * rows + o + r) = v };
3510                            }
3511                            bi += 1;
3512                        }
3513                        o += 2;
3514                    } else {
3515                        let row = &q[o * cols..(o + 1) * cols];
3516                        for (bi, act) in acts.iter().enumerate() {
3517                            let v = row_dot_avx2(row, act) * row_scale[o];
3518                            unsafe { *out_addr.at(bi * rows + o) = v };
3519                        }
3520                        o += 1;
3521                    }
3522                }
3523            };
3524            dispatch_rows(pool, rows, &run);
3525            return;
3526        }
3527        let run = |start: usize, end: usize| {
3528            for o in start..end {
3529                let row = &q[o * cols..(o + 1) * cols];
3530                for (bi, act) in acts.iter().enumerate() {
3531                    let v = row_dot_avx2(row, act) * row_scale[o];
3532                    unsafe { *out_addr.at(bi * rows + o) = v };
3533                }
3534            }
3535        };
3536        dispatch_rows(pool, rows, &run);
3537        return;
3538    }
3539    let out_addr = SendMut(out.as_mut_ptr());
3540    let run = |start: usize, end: usize| {
3541        for o in start..end {
3542            let row = &q[o * cols..(o + 1) * cols];
3543            for (bi, x) in pre.iter().enumerate() {
3544                let mut acc = 0f32;
3545                for j in 0..cols {
3546                    acc += (row[j] as i8) as f32 * x[j];
3547                }
3548                unsafe { *out_addr.at(bi * rows + o) = acc * row_scale[o] };
3549            }
3550        }
3551    };
3552    dispatch_rows(pool, rows, &run);
3553}
3554
3555/// Split rows across pool workers (shared qmatvec pattern). Self-balancing
3556/// — see `Pool::run_rows` for why a static 1/n split is wrong here.
3557fn dispatch_rows(pool: Option<&Pool>, rows: usize, run: &(dyn Fn(usize, usize) + Sync)) {
3558    match pool {
3559        Some(pool) if rows >= 256 => pool.run_rows(rows, run),
3560        _ => run(0, rows),
3561    }
3562}
3563
3564/// Split a q4_block blob into (packed nibbles, f16 group scales).
3565fn q4_split(bytes: &[u8], rows: usize, cols: usize) -> (&[u8], &[u8]) {
3566    let groups = rows * cols / GROUP_SIZE;
3567    bytes.split_at(groups * 16)
3568}
3569
3570/// SIMD unpack for the dominant vbit width B=4 (94% of rows on the
3571/// log2-shape calibration): 16 packed bytes -> 32 centered i8 values.
3572/// vbit packs MSB-first, so the HIGH nibble is the even element
3573/// (opposite of q4_block's lo-first interleave). Centering is u-7.
3574#[inline]
3575fn vbit_fill4(data: &[u8], buf: &mut [u8]) {
3576    #[cfg(target_arch = "aarch64")]
3577    unsafe {
3578        return vbit_fill4_neon(data, buf);
3579    }
3580    #[cfg(target_arch = "x86_64")]
3581    if avx2_enabled() {
3582        return unsafe { vbit_fill4_avx2(data, buf) };
3583    }
3584    #[allow(unreachable_code)]
3585    for (blk, chunk) in buf.chunks_exact_mut(8).enumerate() {
3586        let u = unpack8::<4>(&data[blk * 4..]);
3587        for k in 0..8 {
3588            chunk[k] = (u[k] - 7) as i8 as u8;
3589        }
3590    }
3591}
3592
3593#[cfg(target_arch = "aarch64")]
3594#[target_feature(enable = "neon")]
3595unsafe fn vbit_fill4_neon(data: &[u8], buf: &mut [u8]) {
3596    // SAFETY: buf.len() is a multiple of GROUP_SIZE=32; data holds
3597    // buf.len()/2 packed bytes (validated at load).
3598    unsafe {
3599        use core::arch::aarch64::*;
3600        let n = buf.len();
3601        let mask = vdupq_n_u8(0x0F);
3602        let seven = vdupq_n_s8(7);
3603        let mut g = 0usize;
3604        while g * 32 + 32 <= n {
3605            let b = vld1q_u8(data.as_ptr().add(g * 16));
3606            let hi = vshrq_n_u8::<4>(b);
3607            let lo = vandq_u8(b, mask);
3608            let z0 = vsubq_s8(vreinterpretq_s8_u8(vzip1q_u8(hi, lo)), seven);
3609            let z1 = vsubq_s8(vreinterpretq_s8_u8(vzip2q_u8(hi, lo)), seven);
3610            vst1q_u8(buf.as_mut_ptr().add(g * 32), vreinterpretq_u8_s8(z0));
3611            vst1q_u8(buf.as_mut_ptr().add(g * 32 + 16), vreinterpretq_u8_s8(z1));
3612            g += 1;
3613        }
3614    }
3615}
3616
3617#[cfg(target_arch = "x86_64")]
3618#[target_feature(enable = "avx2")]
3619unsafe fn vbit_fill4_avx2(data: &[u8], buf: &mut [u8]) {
3620    // SAFETY: see vbit_fill4_neon.
3621    unsafe {
3622        use core::arch::x86_64::*;
3623        let n = buf.len();
3624        let mask = _mm_set1_epi8(0x0F);
3625        let seven = _mm256_set1_epi8(7);
3626        let mut g = 0usize;
3627        while g * 32 + 32 <= n {
3628            let b = _mm_loadu_si128(data.as_ptr().add(g * 16) as *const __m128i);
3629            let hi = _mm_and_si128(_mm_srli_epi16::<4>(b), mask);
3630            let lo = _mm_and_si128(b, mask);
3631            let z = _mm256_sub_epi8(
3632                _mm256_set_m128i(_mm_unpackhi_epi8(hi, lo), _mm_unpacklo_epi8(hi, lo)),
3633                seven,
3634            );
3635            _mm256_storeu_si256(buf.as_mut_ptr().add(g * 32) as *mut __m256i, z);
3636            g += 1;
3637        }
3638    }
3639}
3640
3641/// Unpack 8 MSB-first B-bit values from exactly B bytes (fixed shifts —
3642/// no serial bit-buffer, auto-vectorizable). Every 32-value group starts
3643/// byte-aligned (32·B/8 is integral for B∈3..8), so groups decompose
3644/// into 4 such blocks.
3645#[inline(always)]
3646fn unpack8<const B: usize>(data: &[u8]) -> [i32; 8] {
3647    let mut acc = 0u64;
3648    for i in 0..B {
3649        acc = (acc << 8) | data[i] as u64;
3650    }
3651    let mask = (1u64 << B) - 1;
3652    let mut out = [0i32; 8];
3653    for (k, o) in out.iter_mut().enumerate() {
3654        *o = ((acc >> ((7 - k) * B)) & mask) as i32;
3655    }
3656    out
3657}
3658
3659/// Fused vbit matvec straight from the mapped bytes (spec §3, P13
3660/// FIG.3): [u8 bits: rows][f16 scales: rows·cols/32][bit-packed rows,
3661/// MSB-first, byte-padded]. Row data offsets are precomputed at load
3662/// (`vbit_row_offsets`) — the per-call prefix scan was O(rows) pure
3663/// overhead on every matvec.
3664#[allow(clippy::too_many_arguments)]
3665fn vbitmatvec(
3666    bytes: &[u8],
3667    offsets: &[usize],
3668    x: &[f32],
3669    rows: usize,
3670    cols: usize,
3671    out: &mut [f32],
3672    pool: Option<&Pool>,
3673) {
3674    debug_assert_eq!(out.len(), rows);
3675    debug_assert_eq!(offsets.len(), rows + 1);
3676
3677    // SDOT path: unpack the row to centered i8 once, then per-group
3678    // int8 dot against the quantized activations — same A8W8 contract
3679    // as q8 (bounded noise; CMF_SDOT=0 keeps the exact scalar path).
3680    if a8w8_enabled() {
3681        let act = split_act(x);
3682        let out_addr = SendMut(out.as_mut_ptr());
3683        let run = move |start: usize, end: usize| {
3684            vbit_range_a8w8(bytes, offsets, x, &act, rows, cols, out_addr, start, end)
3685        };
3686        dispatch_rows(pool, rows, &run);
3687        return;
3688    }
3689
3690    let out_addr = SendMut(out.as_mut_ptr());
3691    let run = move |start: usize, end: usize| {
3692        vbit_range_f32(bytes, offsets, x, rows, cols, out_addr, start, end)
3693    };
3694    dispatch_rows(pool, rows, &run);
3695}
3696
3697/// One vbit row range via the A8W8 int8 path — kernel body of
3698/// `vbitmatvec`, extracted so multi-matrix jobs can drive it for
3699/// several tensors in one dispatch (b=8 rows go exact f32).
3700#[allow(clippy::too_many_arguments)]
3701fn vbit_range_a8w8(
3702    bytes: &[u8],
3703    offsets: &[usize],
3704    x: &[f32],
3705    act: &SplitAct,
3706    rows: usize,
3707    cols: usize,
3708    out: SendMut,
3709    start: usize,
3710    end: usize,
3711) {
3712    let ng = cols / GROUP_SIZE;
3713    let bits = &bytes[..rows];
3714    let sc_off = rows;
3715    let row_dot = |r: usize| -> f32 {
3716        let b = bits[r] as usize;
3717        let l = (1i32 << (b - 1)) - 1;
3718        let mask = (1u64 << b) - 1;
3719        let data = &bytes[offsets[r]..offsets[r + 1]];
3720        if b == 8 {
3721            // u−L reaches 128 → does not fit i8; exact f32 path.
3722            let (mut acc, mut nbits, mut idx) = (0u64, 0usize, 0usize);
3723            let mut dot = 0f32;
3724            for g in 0..ng {
3725                let so = (r * ng + g) * 2;
3726                let sgf = f16_to_f32(u16::from_le_bytes([
3727                    bytes[sc_off + so],
3728                    bytes[sc_off + so + 1],
3729                ]));
3730                let xg = &x[g * GROUP_SIZE..(g + 1) * GROUP_SIZE];
3731                let mut gd = 0f32;
3732                for &xv in xg.iter() {
3733                    if nbits < 8 {
3734                        acc = (acc << 8) | data[idx] as u64;
3735                        idx += 1;
3736                        nbits += 8;
3737                    }
3738                    let u = ((acc >> (nbits - 8)) & 0xFF) as i32;
3739                    nbits -= 8;
3740                    gd += (u - l) as f32 * xv;
3741                }
3742                dot += gd * sgf;
3743            }
3744            return dot;
3745        }
3746        // Per-worker scratch: this closure runs for every row of the
3747        // tensor (lm_head ≈ 150k rows/token) — a heap allocation per
3748        // row was measurable pure overhead.
3749        thread_local! {
3750            static VBIT_SCRATCH: std::cell::RefCell<Vec<u8>> =
3751                const { std::cell::RefCell::new(Vec::new()) };
3752        }
3753        #[inline(always)]
3754        fn fill<const B: usize>(data: &[u8], l: i32, buf: &mut [u8]) {
3755            for (blk, chunk) in buf.chunks_exact_mut(8).enumerate() {
3756                let u = unpack8::<B>(&data[blk * B..]);
3757                for k in 0..8 {
3758                    chunk[k] = (u[k] - l) as i8 as u8;
3759                }
3760            }
3761        }
3762        let _ = mask;
3763        VBIT_SCRATCH.with(|scratch| {
3764            let mut buf = scratch.borrow_mut();
3765            buf.resize(cols, 0);
3766            match b {
3767                3 => fill::<3>(data, l, &mut buf),
3768                4 => vbit_fill4(data, &mut buf),
3769                5 => fill::<5>(data, l, &mut buf),
3770                6 => fill::<6>(data, l, &mut buf),
3771                _ => unreachable!(),
3772            }
3773            let mut dot = 0f32;
3774            for g in 0..ng {
3775                let so = (r * ng + g) * 2;
3776                let s = f16_to_f32(u16::from_le_bytes([
3777                    bytes[sc_off + so],
3778                    bytes[sc_off + so + 1],
3779                ]));
3780                let d = dot_i8_i8(
3781                    &buf[g * GROUP_SIZE..(g + 1) * GROUP_SIZE],
3782                    &act.xq[g * GROUP_SIZE..(g + 1) * GROUP_SIZE],
3783                ) as f32
3784                    * act.sx;
3785                dot += d * s;
3786            }
3787            for &(j, xv) in &act.outliers {
3788                let so = (r * ng + j / GROUP_SIZE) * 2;
3789                let s = f16_to_f32(u16::from_le_bytes([
3790                    bytes[sc_off + so],
3791                    bytes[sc_off + so + 1],
3792                ]));
3793                // xq is zeroed at outlier slots — add the exact term.
3794                dot += (buf[j] as i8) as f32 * s * xv;
3795            }
3796            dot
3797        })
3798    };
3799    for r in start..end {
3800        // SAFETY: disjoint row ranges per worker.
3801        unsafe { *out.at(r) = row_dot(r) };
3802    }
3803}
3804
3805/// Exact scalar vbit row range (same extraction, non-SDOT path).
3806#[allow(clippy::too_many_arguments)]
3807fn vbit_range_f32(
3808    bytes: &[u8],
3809    offsets: &[usize],
3810    x: &[f32],
3811    rows: usize,
3812    cols: usize,
3813    out: 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    // Per-bit-width specialized inner loops: the compiler unrolls the
3821    // constant shifts (the generic bit-buffer loop was branch-bound —
3822    // 5.6 vs 13.2 tok/s q4 on the 0.8B).
3823    #[inline(always)]
3824    fn dot_row<const B: usize>(
3825        data: &[u8],
3826        bytes: &[u8],
3827        sc_off: usize,
3828        r: usize,
3829        ng: usize,
3830        x: &[f32],
3831    ) -> f32 {
3832        let l = ((1i32 << (B - 1)) - 1) as f32;
3833        let gbytes = GROUP_SIZE * B / 8;
3834        let mut dot = 0f32;
3835        for g in 0..ng {
3836            let so = (r * ng + g) * 2;
3837            let s = f16_to_f32(u16::from_le_bytes([
3838                bytes[sc_off + so],
3839                bytes[sc_off + so + 1],
3840            ]));
3841            let xg = &x[g * GROUP_SIZE..(g + 1) * GROUP_SIZE];
3842            let gd0 = &data[g * gbytes..(g + 1) * gbytes];
3843            let mut gd = 0f32;
3844            for blk in 0..GROUP_SIZE / 8 {
3845                let u = unpack8::<B>(&gd0[blk * B..]);
3846                let xb = &xg[blk * 8..blk * 8 + 8];
3847                for k in 0..8 {
3848                    gd += (u[k] as f32 - l) * xb[k];
3849                }
3850            }
3851            dot += gd * s;
3852        }
3853        dot
3854    }
3855    for r in start..end {
3856        let data = &bytes[offsets[r]..offsets[r + 1]];
3857        let v = match bits[r] {
3858            3 => dot_row::<3>(data, bytes, sc_off, r, ng, x),
3859            4 => dot_row::<4>(data, bytes, sc_off, r, ng, x),
3860            5 => dot_row::<5>(data, bytes, sc_off, r, ng, x),
3861            6 => dot_row::<6>(data, bytes, sc_off, r, ng, x),
3862            8 => dot_row::<8>(data, bytes, sc_off, r, ng, x),
3863            b => unreachable!("vbit bit-width {b} (validated at load)"),
3864        };
3865        // SAFETY: disjoint row ranges per worker.
3866        unsafe { *out.at(r) = v };
3867    }
3868}
3869
3870/// Fused two-input vbit matvec: each row is unpacked from the mmap ONCE
3871/// and dotted against BOTH activations (MTP verify / pair prefill used
3872/// to run two full matvecs — double weight traffic and double unpack).
3873/// Per-input math is identical to `vbitmatvec` → same accuracy contract.
3874#[allow(clippy::too_many_arguments)]
3875fn vbitmatvec2(
3876    bytes: &[u8],
3877    offsets: &[usize],
3878    x1: &[f32],
3879    x2: &[f32],
3880    rows: usize,
3881    cols: usize,
3882    o1: &mut [f32],
3883    o2: &mut [f32],
3884    pool: Option<&Pool>,
3885) {
3886    debug_assert_eq!(o1.len(), rows);
3887    debug_assert_eq!(o2.len(), rows);
3888
3889    if a8w8_enabled() {
3890        let a1 = split_act(x1);
3891        let a2 = split_act(x2);
3892        let p1 = SendMut(o1.as_mut_ptr());
3893        let p2 = SendMut(o2.as_mut_ptr());
3894        let run = move |start: usize, end: usize| {
3895            vbit_range2_a8w8(
3896                bytes, offsets, x1, x2, &a1, &a2, rows, cols, p1, p2, start, end,
3897            )
3898        };
3899        dispatch_rows(pool, rows, &run);
3900        return;
3901    }
3902
3903    let p1 = SendMut(o1.as_mut_ptr());
3904    let p2 = SendMut(o2.as_mut_ptr());
3905    let run = move |start: usize, end: usize| {
3906        vbit_range2_f32(bytes, offsets, x1, x2, rows, cols, p1, p2, start, end)
3907    };
3908    dispatch_rows(pool, rows, &run);
3909}
3910
3911/// Two-input vbit row range via the A8W8 int8 path — kernel body of
3912/// `vbitmatvec2`, extracted for pair multi-matrix jobs (b=8 rows go
3913/// exact f32 for both lanes, bits streamed once).
3914#[allow(clippy::too_many_arguments)]
3915fn vbit_range2_a8w8(
3916    bytes: &[u8],
3917    offsets: &[usize],
3918    x1: &[f32],
3919    x2: &[f32],
3920    a1: &SplitAct,
3921    a2: &SplitAct,
3922    rows: usize,
3923    cols: usize,
3924    p1: SendMut,
3925    p2: SendMut,
3926    start: usize,
3927    end: usize,
3928) {
3929    let ng = cols / GROUP_SIZE;
3930    let bits = &bytes[..rows];
3931    let sc_off = rows;
3932    let row_dots = |r: usize| -> (f32, f32) {
3933        let b = bits[r] as usize;
3934        let l = (1i32 << (b - 1)) - 1;
3935        let data = &bytes[offsets[r]..offsets[r + 1]];
3936        if b == 8 {
3937            // u−L reaches 128 → does not fit i8; exact f32 path,
3938            // bits still streamed once for both lanes.
3939            let (mut acc, mut nbits, mut idx) = (0u64, 0usize, 0usize);
3940            let (mut d1, mut d2) = (0f32, 0f32);
3941            for g in 0..ng {
3942                let so = (r * ng + g) * 2;
3943                let sgf = f16_to_f32(u16::from_le_bytes([
3944                    bytes[sc_off + so],
3945                    bytes[sc_off + so + 1],
3946                ]));
3947                let (mut g1, mut g2) = (0f32, 0f32);
3948                for k in 0..GROUP_SIZE {
3949                    if nbits < 8 {
3950                        acc = (acc << 8) | data[idx] as u64;
3951                        idx += 1;
3952                        nbits += 8;
3953                    }
3954                    let u = ((acc >> (nbits - 8)) & 0xFF) as i32;
3955                    nbits -= 8;
3956                    let w = (u - l) as f32;
3957                    g1 += w * x1[g * GROUP_SIZE + k];
3958                    g2 += w * x2[g * GROUP_SIZE + k];
3959                }
3960                d1 += g1 * sgf;
3961                d2 += g2 * sgf;
3962            }
3963            return (d1, d2);
3964        }
3965        thread_local! {
3966            static VBIT_SCRATCH2: std::cell::RefCell<Vec<u8>> =
3967                const { std::cell::RefCell::new(Vec::new()) };
3968        }
3969        #[inline(always)]
3970        fn fill<const B: usize>(data: &[u8], l: i32, buf: &mut [u8]) {
3971            for (blk, chunk) in buf.chunks_exact_mut(8).enumerate() {
3972                let u = unpack8::<B>(&data[blk * B..]);
3973                for k in 0..8 {
3974                    chunk[k] = (u[k] - l) as i8 as u8;
3975                }
3976            }
3977        }
3978        VBIT_SCRATCH2.with(|scratch| {
3979            let mut buf = scratch.borrow_mut();
3980            buf.resize(cols, 0);
3981            match b {
3982                3 => fill::<3>(data, l, &mut buf),
3983                4 => vbit_fill4(data, &mut buf),
3984                5 => fill::<5>(data, l, &mut buf),
3985                6 => fill::<6>(data, l, &mut buf),
3986                _ => unreachable!(),
3987            }
3988            let (mut d1, mut d2) = (0f32, 0f32);
3989            for g in 0..ng {
3990                let so = (r * ng + g) * 2;
3991                let s = f16_to_f32(u16::from_le_bytes([
3992                    bytes[sc_off + so],
3993                    bytes[sc_off + so + 1],
3994                ]));
3995                let wg = &buf[g * GROUP_SIZE..(g + 1) * GROUP_SIZE];
3996                let v1 = dot_i8_i8(wg, &a1.xq[g * GROUP_SIZE..(g + 1) * GROUP_SIZE]) as f32 * a1.sx;
3997                let v2 = dot_i8_i8(wg, &a2.xq[g * GROUP_SIZE..(g + 1) * GROUP_SIZE]) as f32 * a2.sx;
3998                d1 += v1 * s;
3999                d2 += v2 * s;
4000            }
4001            for &(j, xv) in &a1.outliers {
4002                let so = (r * ng + j / GROUP_SIZE) * 2;
4003                let s = f16_to_f32(u16::from_le_bytes([
4004                    bytes[sc_off + so],
4005                    bytes[sc_off + so + 1],
4006                ]));
4007                d1 += (buf[j] as i8) as f32 * s * xv;
4008            }
4009            for &(j, xv) in &a2.outliers {
4010                let so = (r * ng + j / GROUP_SIZE) * 2;
4011                let s = f16_to_f32(u16::from_le_bytes([
4012                    bytes[sc_off + so],
4013                    bytes[sc_off + so + 1],
4014                ]));
4015                d2 += (buf[j] as i8) as f32 * s * xv;
4016            }
4017            (d1, d2)
4018        })
4019    };
4020    for r in start..end {
4021        let (v1, v2) = row_dots(r);
4022        // SAFETY: disjoint row ranges per worker.
4023        unsafe {
4024            *p1.at(r) = v1;
4025            *p2.at(r) = v2;
4026        }
4027    }
4028}
4029
4030/// Two-input exact scalar vbit row range (same extraction) —
4031/// per-bit-width specialized, two accumulators per row; per-lane
4032/// accumulation order matches `vbitmatvec` exactly.
4033#[allow(clippy::too_many_arguments)]
4034fn vbit_range2_f32(
4035    bytes: &[u8],
4036    offsets: &[usize],
4037    x1: &[f32],
4038    x2: &[f32],
4039    rows: usize,
4040    cols: usize,
4041    p1: SendMut,
4042    p2: SendMut,
4043    start: usize,
4044    end: usize,
4045) {
4046    let ng = cols / GROUP_SIZE;
4047    let bits = &bytes[..rows];
4048    let sc_off = rows;
4049    #[inline(always)]
4050    #[allow(clippy::too_many_arguments)]
4051    fn dot_row2<const B: usize>(
4052        data: &[u8],
4053        bytes: &[u8],
4054        sc_off: usize,
4055        r: usize,
4056        ng: usize,
4057        x1: &[f32],
4058        x2: &[f32],
4059    ) -> (f32, f32) {
4060        let l = ((1i32 << (B - 1)) - 1) as f32;
4061        let gbytes = GROUP_SIZE * B / 8;
4062        let (mut d1, mut d2) = (0f32, 0f32);
4063        for g in 0..ng {
4064            let so = (r * ng + g) * 2;
4065            let s = f16_to_f32(u16::from_le_bytes([
4066                bytes[sc_off + so],
4067                bytes[sc_off + so + 1],
4068            ]));
4069            let x1g = &x1[g * GROUP_SIZE..(g + 1) * GROUP_SIZE];
4070            let x2g = &x2[g * GROUP_SIZE..(g + 1) * GROUP_SIZE];
4071            let gd0 = &data[g * gbytes..(g + 1) * gbytes];
4072            let (mut g1, mut g2) = (0f32, 0f32);
4073            for blk in 0..GROUP_SIZE / 8 {
4074                let u = unpack8::<B>(&gd0[blk * B..]);
4075                for k in 0..8 {
4076                    let w = u[k] as f32 - l;
4077                    g1 += w * x1g[blk * 8 + k];
4078                    g2 += w * x2g[blk * 8 + k];
4079                }
4080            }
4081            d1 += g1 * s;
4082            d2 += g2 * s;
4083        }
4084        (d1, d2)
4085    }
4086    for r in start..end {
4087        let data = &bytes[offsets[r]..offsets[r + 1]];
4088        let (v1, v2) = match bits[r] {
4089            3 => dot_row2::<3>(data, bytes, sc_off, r, ng, x1, x2),
4090            4 => dot_row2::<4>(data, bytes, sc_off, r, ng, x1, x2),
4091            5 => dot_row2::<5>(data, bytes, sc_off, r, ng, x1, x2),
4092            6 => dot_row2::<6>(data, bytes, sc_off, r, ng, x1, x2),
4093            8 => dot_row2::<8>(data, bytes, sc_off, r, ng, x1, x2),
4094            b => unreachable!("vbit bit-width {b} (validated at load)"),
4095        };
4096        // SAFETY: disjoint row ranges per worker.
4097        unsafe {
4098            *p1.at(r) = v1;
4099            *p2.at(r) = v2;
4100        }
4101    }
4102}
4103
4104// ───────────────────── q4_tiled kernels (§4.3) ─────────────────────
4105
4106/// One q4_tiled row dot on the A8W8 int8 path: per 32-group the tile
4107/// is ONE sequential read — [f16 scale][16B nibbles] — versus the two
4108/// distant streams of the split layout. Values/order identical to the
4109/// split kernels.
4110#[inline]
4111#[allow(unreachable_code)]
4112fn dot_q4t_row_i8(bytes: &[u8], r: usize, gpr: usize, xq: &[i8]) -> f32 {
4113    #[cfg(target_arch = "aarch64")]
4114    unsafe {
4115        return dot_q4t_row_sdot(bytes, r, gpr, xq);
4116    }
4117    #[cfg(target_arch = "x86_64")]
4118    unsafe {
4119        if vnni_tiles_enabled() {
4120            return dot_q4t_row_vnni(bytes, r, gpr, xq);
4121        }
4122        return dot_q4t_row_avx2(bytes, r, gpr, xq);
4123    }
4124    let mut acc = 0f32;
4125    for gi in 0..gpr {
4126        let tile = &bytes[(r * gpr + gi) * Q4_TILE..(r * gpr + gi + 1) * Q4_TILE];
4127        let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
4128        let mut d = 0i32;
4129        for (k, &b) in tile[2..].iter().enumerate() {
4130            d += ((b & 0x0F) as i32 - 8) * xq[gi * GROUP_SIZE + k * 2] as i32
4131                + (((b >> 4) & 0x0F) as i32 - 8) * xq[gi * GROUP_SIZE + k * 2 + 1] as i32;
4132        }
4133        acc += d as f32 * s;
4134    }
4135    acc
4136}
4137
4138#[cfg(target_arch = "aarch64")]
4139#[target_feature(enable = "neon,dotprod")]
4140unsafe fn dot_q4t_row_sdot(bytes: &[u8], r: usize, gpr: usize, xq: &[i8]) -> f32 {
4141    // SAFETY: callers uphold slice-length contracts (18B tile per group,
4142    // xq.len() == gpr·GROUP_SIZE).
4143    unsafe {
4144        use core::arch::aarch64::*;
4145        use core::arch::asm;
4146        let lomask = vdupq_n_u8(0x0F);
4147        let eight = vdupq_n_s8(8);
4148        let mut acc = 0f32;
4149        for gi in 0..gpr {
4150            let t = bytes.as_ptr().add((r * gpr + gi) * Q4_TILE);
4151            let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
4152            let b = vld1q_u8(t.add(2));
4153            let lo = vandq_u8(b, lomask);
4154            let hi = vshrq_n_u8::<4>(b);
4155            let e0 = vsubq_s8(vreinterpretq_s8_u8(vzip1q_u8(lo, hi)), eight);
4156            let e1 = vsubq_s8(vreinterpretq_s8_u8(vzip2q_u8(lo, hi)), eight);
4157            let x0 = vld1q_s8(xq.as_ptr().add(gi * GROUP_SIZE));
4158            let x1 = vld1q_s8(xq.as_ptr().add(gi * GROUP_SIZE + 16));
4159            let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
4160            asm!(
4161                "sdot {a0:v}.4s, {e0:v}.16b, {x0:v}.16b",
4162                "sdot {a1:v}.4s, {e1:v}.16b, {x1:v}.16b",
4163                a0 = inout(vreg) a0, a1 = inout(vreg) a1,
4164                e0 = in(vreg) e0, x0 = in(vreg) x0, e1 = in(vreg) e1, x1 = in(vreg) x1,
4165                options(pure, nomem, nostack),
4166            );
4167            acc += vaddvq_s32(vaddq_s32(a0, a1)) as f32 * s;
4168        }
4169        acc
4170    }
4171}
4172
4173#[cfg(target_arch = "x86_64")]
4174#[target_feature(enable = "avx2")]
4175unsafe fn dot_q4t_row_avx2(bytes: &[u8], r: usize, gpr: usize, xq: &[i8]) -> f32 {
4176    // SAFETY: see dot_q4t_row_sdot.
4177    unsafe {
4178        use core::arch::x86_64::*;
4179        let lomask = _mm_set1_epi8(0x0F);
4180        let eight = _mm256_set1_epi8(8);
4181        let ones = _mm256_set1_epi16(1);
4182        let mut acc = 0f32;
4183        for gi in 0..gpr {
4184            let t = bytes.as_ptr().add((r * gpr + gi) * Q4_TILE);
4185            let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
4186            let b = _mm_loadu_si128(t.add(2) as *const __m128i);
4187            let lo = _mm_and_si128(b, lomask);
4188            let hi = _mm_and_si128(_mm_srli_epi16::<4>(b), lomask);
4189            let w = _mm256_sub_epi8(
4190                _mm256_set_m128i(_mm_unpackhi_epi8(lo, hi), _mm_unpacklo_epi8(lo, hi)),
4191                eight,
4192            );
4193            let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
4194            let p16 = _mm256_maddubs_epi16(_mm256_abs_epi8(w), _mm256_sign_epi8(x, w));
4195            let d = _mm256_madd_epi16(p16, ones);
4196            let hi128 = _mm256_extracti128_si256::<1>(d);
4197            let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
4198            let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
4199            let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
4200            acc += _mm_cvtsi128_si32(s32) as f32 * s;
4201        }
4202        acc
4203    }
4204}
4205
4206/// VNNI twin of `dot_q4t_row_avx2`: same unpack, `vpdpbusd` replaces
4207/// the maddubs+madd pair (see `dpbusd_hsum` — sums are bit-identical).
4208/// 256-bit VL encoding, so the VEX `vpsignb` stays usable.
4209#[cfg(target_arch = "x86_64")]
4210#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
4211unsafe fn dot_q4t_row_vnni(bytes: &[u8], r: usize, gpr: usize, xq: &[i8]) -> f32 {
4212    // SAFETY: see dot_q4t_row_sdot.
4213    unsafe {
4214        use core::arch::x86_64::*;
4215        let lomask = _mm_set1_epi8(0x0F);
4216        let eight = _mm256_set1_epi8(8);
4217        let mut acc = 0f32;
4218        for gi in 0..gpr {
4219            let t = bytes.as_ptr().add((r * gpr + gi) * Q4_TILE);
4220            let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
4221            let b = _mm_loadu_si128(t.add(2) as *const __m128i);
4222            let lo = _mm_and_si128(b, lomask);
4223            let hi = _mm_and_si128(_mm_srli_epi16::<4>(b), lomask);
4224            let w = _mm256_sub_epi8(
4225                _mm256_set_m128i(_mm_unpackhi_epi8(lo, hi), _mm_unpacklo_epi8(lo, hi)),
4226                eight,
4227            );
4228            let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
4229            let d = dpbusd_hsum(_mm256_abs_epi8(w), _mm256_sign_epi8(x, w));
4230            acc += d as f32 * s;
4231        }
4232        acc
4233    }
4234}
4235
4236/// One q4_tiled row against FOUR activation streams: the nibble unpack
4237/// and abs() happen once per group instead of once per (group,
4238/// activation) — the unpack is the dominant per-element cost of the
4239/// tiled format (roadmap P0 portable blocking, q4t leg).
4240#[cfg(target_arch = "x86_64")]
4241// `fma` is NOT implied by `avx2`: without it LLVM lowers _mm256_fmadd_ps
4242// to a libm call per lane — measured 2x slower than the reduction this
4243// kernel replaces. The runtime gate (`avx2_enabled`) already requires
4244// both features, so declaring it here is safe.
4245#[target_feature(enable = "avx2,fma")]
4246unsafe fn dot_q4t_row_1x4_avx2(bytes: &[u8], r: usize, gpr: usize, xs: [&[i8]; 4]) -> [f32; 4] {
4247    // SAFETY: callers uphold the 18B-tile and xq-length contracts.
4248    unsafe {
4249        use core::arch::x86_64::*;
4250        let lomask = _mm_set1_epi8(0x0F);
4251        let eight = _mm256_set1_epi8(8);
4252        let ones = _mm256_set1_epi16(1);
4253        // One f32 accumulator VECTOR per activation, reduced once at the
4254        // end. Folding each group's i32 lanes to a scalar inside the loop
4255        // costs an extracti128 + three shift/add + a movd — a cross-lane
4256        // dependency chain per (group, activation), 288 of them per row at
4257        // cols=2304. The per-group scale is what forces a float
4258        // accumulator; it does not force a horizontal sum.
4259        //
4260        // The four accumulators are NAMED, not an array: as `[__m256; 4]`
4261        // indexed by a loop variable LLVM keeps them in memory and every
4262        // group pays four 32-byte loads and stores. That alone made this
4263        // kernel 2x SLOWER than the per-group reduction it replaces
4264        // (measured on the EPYC box: 150 s vs 71 s for two 256² steps).
4265        let mut f0 = _mm256_setzero_ps();
4266        let mut f1 = _mm256_setzero_ps();
4267        let mut f2 = _mm256_setzero_ps();
4268        let mut f3 = _mm256_setzero_ps();
4269        for gi in 0..gpr {
4270            let t = bytes.as_ptr().add((r * gpr + gi) * Q4_TILE);
4271            let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
4272            let sv = _mm256_set1_ps(s);
4273            let bb = _mm_loadu_si128(t.add(2) as *const __m128i);
4274            let lo = _mm_and_si128(bb, lomask);
4275            let hi = _mm_and_si128(_mm_srli_epi16::<4>(bb), lomask);
4276            let w = _mm256_sub_epi8(
4277                _mm256_set_m128i(_mm_unpackhi_epi8(lo, hi), _mm_unpacklo_epi8(lo, hi)),
4278                eight,
4279            );
4280            let aw = _mm256_abs_epi8(w);
4281            let off = gi * GROUP_SIZE;
4282            let dot = |xq: &[i8]| {
4283                let x = _mm256_loadu_si256(xq.as_ptr().add(off) as *const __m256i);
4284                let p16 = _mm256_maddubs_epi16(aw, _mm256_sign_epi8(x, w));
4285                _mm256_cvtepi32_ps(_mm256_madd_epi16(p16, ones))
4286            };
4287            f0 = _mm256_fmadd_ps(dot(xs[0]), sv, f0);
4288            f1 = _mm256_fmadd_ps(dot(xs[1]), sv, f1);
4289            f2 = _mm256_fmadd_ps(dot(xs[2]), sv, f2);
4290            f3 = _mm256_fmadd_ps(dot(xs[3]), sv, f3);
4291        }
4292        [
4293            hsum256_ps(f0),
4294            hsum256_ps(f1),
4295            hsum256_ps(f2),
4296            hsum256_ps(f3),
4297        ]
4298    }
4299}
4300
4301/// Horizontal sum of eight f32 lanes — the one cross-lane reduction the
4302/// blocked kernels pay, once per row instead of once per group.
4303#[cfg(target_arch = "x86_64")]
4304#[target_feature(enable = "avx2")]
4305#[inline]
4306unsafe fn hsum256_ps(v: core::arch::x86_64::__m256) -> f32 {
4307    // SAFETY: pure register arithmetic on the caller's vector.
4308    unsafe {
4309        use core::arch::x86_64::*;
4310        let hi = _mm256_extractf128_ps::<1>(v);
4311        let s = _mm_add_ps(_mm256_castps256_ps128(v), hi);
4312        let s = _mm_add_ps(s, _mm_movehl_ps(s, s));
4313        let s = _mm_add_ss(s, _mm_shuffle_ps::<0x55>(s, s));
4314        _mm_cvtss_f32(s)
4315    }
4316}
4317
4318/// VNNI twin of `dot_q4t_row_1x4_avx2` (see `dpbusd_hsum`).
4319#[cfg(target_arch = "x86_64")]
4320#[target_feature(enable = "avx2,fma,avx512f,avx512bw,avx512vl,avx512vnni")]
4321unsafe fn dot_q4t_row_1x4_vnni(bytes: &[u8], r: usize, gpr: usize, xs: [&[i8]; 4]) -> [f32; 4] {
4322    // SAFETY: callers uphold the 18B-tile and xq-length contracts.
4323    unsafe {
4324        use core::arch::x86_64::*;
4325        let lomask = _mm_set1_epi8(0x0F);
4326        let eight = _mm256_set1_epi8(8);
4327        // Same shape as the AVX2 twin: accumulate in f32 vectors and pay
4328        // one cross-lane reduction per row, not per (group, activation).
4329        let mut f0 = _mm256_setzero_ps();
4330        let mut f1 = _mm256_setzero_ps();
4331        let mut f2 = _mm256_setzero_ps();
4332        let mut f3 = _mm256_setzero_ps();
4333        for gi in 0..gpr {
4334            let t = bytes.as_ptr().add((r * gpr + gi) * Q4_TILE);
4335            let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
4336            let sv = _mm256_set1_ps(s);
4337            let bb = _mm_loadu_si128(t.add(2) as *const __m128i);
4338            let lo = _mm_and_si128(bb, lomask);
4339            let hi = _mm_and_si128(_mm_srli_epi16::<4>(bb), lomask);
4340            let w = _mm256_sub_epi8(
4341                _mm256_set_m128i(_mm_unpackhi_epi8(lo, hi), _mm_unpacklo_epi8(lo, hi)),
4342                eight,
4343            );
4344            let aw = _mm256_abs_epi8(w);
4345            let off = gi * GROUP_SIZE;
4346            let dot = |xq: &[i8]| {
4347                let x = _mm256_loadu_si256(xq.as_ptr().add(off) as *const __m256i);
4348                _mm256_cvtepi32_ps(_mm256_dpbusd_epi32(
4349                    _mm256_setzero_si256(),
4350                    aw,
4351                    _mm256_sign_epi8(x, w),
4352                ))
4353            };
4354            f0 = _mm256_fmadd_ps(dot(xs[0]), sv, f0);
4355            f1 = _mm256_fmadd_ps(dot(xs[1]), sv, f1);
4356            f2 = _mm256_fmadd_ps(dot(xs[2]), sv, f2);
4357            f3 = _mm256_fmadd_ps(dot(xs[3]), sv, f3);
4358        }
4359        let acc = [
4360            hsum256_ps(f0),
4361            hsum256_ps(f1),
4362            hsum256_ps(f2),
4363            hsum256_ps(f3),
4364        ];
4365        acc
4366    }
4367}
4368
4369/// ARM twin of `dot_q4t_row_1x4_avx2`: one nibble unpack per group
4370/// serves FOUR activation streams. Per stream the group order and f32
4371/// accumulation match `dot_q4t_row_sdot` exactly — batch == matvec
4372/// bit-for-bit.
4373#[cfg(target_arch = "aarch64")]
4374#[target_feature(enable = "neon,dotprod")]
4375unsafe fn dot_q4t_row_1x4_sdot(bytes: &[u8], r: usize, gpr: usize, xs: [&[i8]; 4]) -> [f32; 4] {
4376    // SAFETY: callers uphold the 18B-tile and xq-length contracts.
4377    unsafe {
4378        use core::arch::aarch64::*;
4379        use core::arch::asm;
4380        let lomask = vdupq_n_u8(0x0F);
4381        let eight = vdupq_n_s8(8);
4382        let mut acc = [0f32; 4];
4383        for gi in 0..gpr {
4384            let t = bytes.as_ptr().add((r * gpr + gi) * Q4_TILE);
4385            let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
4386            let b = vld1q_u8(t.add(2));
4387            let lo = vandq_u8(b, lomask);
4388            let hi = vshrq_n_u8::<4>(b);
4389            let e0 = vsubq_s8(vreinterpretq_s8_u8(vzip1q_u8(lo, hi)), eight);
4390            let e1 = vsubq_s8(vreinterpretq_s8_u8(vzip2q_u8(lo, hi)), eight);
4391            for (k, xq) in xs.iter().enumerate() {
4392                let x0 = vld1q_s8(xq.as_ptr().add(gi * GROUP_SIZE));
4393                let x1 = vld1q_s8(xq.as_ptr().add(gi * GROUP_SIZE + 16));
4394                let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
4395                asm!(
4396                    "sdot {a0:v}.4s, {e0:v}.16b, {x0:v}.16b",
4397                    "sdot {a1:v}.4s, {e1:v}.16b, {x1:v}.16b",
4398                    a0 = inout(vreg) a0, a1 = inout(vreg) a1,
4399                    e0 = in(vreg) e0, x0 = in(vreg) x0, e1 = in(vreg) e1, x1 = in(vreg) x1,
4400                    options(pure, nomem, nostack),
4401                );
4402                acc[k] += vaddvq_s32(vaddq_s32(a0, a1)) as f32 * s;
4403            }
4404        }
4405        acc
4406    }
4407}
4408
4409/// Exact-term correction for A8W8 outliers on a tiled row.
4410#[inline]
4411fn q4t_outlier(bytes: &[u8], r: usize, gpr: usize, j: usize) -> (f32, f32) {
4412    let gi = j / GROUP_SIZE;
4413    let k = j % GROUP_SIZE;
4414    let tile = &bytes[(r * gpr + gi) * Q4_TILE..(r * gpr + gi + 1) * Q4_TILE];
4415    let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
4416    let byte = tile[2 + k / 2];
4417    let nib = if k & 1 == 0 { byte & 0x0F } else { byte >> 4 };
4418    ((nib as i32 - 8) as f32, s)
4419}
4420
4421/// Exact scalar q4_tiled row (CMF_SDOT=0 contract) — same pairwise
4422/// accumulation shape as `q4_range_f32`.
4423#[inline]
4424fn q4t_row_exact(bytes: &[u8], r: usize, gpr: usize, x: &[f32]) -> f32 {
4425    let mut acc = 0f32;
4426    for gi in 0..gpr {
4427        let tile = &bytes[(r * gpr + gi) * Q4_TILE..(r * gpr + gi + 1) * Q4_TILE];
4428        let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
4429        let xg = &x[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
4430        let mut ga = 0f32;
4431        for (k, &b) in tile[2..].iter().enumerate() {
4432            ga += ((b & 0x0F) as f32 - 8.0) * xg[k * 2]
4433                + (((b >> 4) & 0x0F) as f32 - 8.0) * xg[k * 2 + 1];
4434        }
4435        acc += ga * s;
4436    }
4437    acc
4438}
4439
4440/// Split view of a `q4tp` payload. The three planes are resolved once per
4441/// matvec instead of per row — `q4tp_sections` is cheap, but doing it inside
4442/// the row loop would put a division on the hot path for nothing.
4443struct Q4tpView<'a> {
4444    nib: &'a [u8],
4445    params: &'a [u8],
4446    codes: &'a [u8],
4447    stride: usize,
4448    /// q2tp reads the ladder with rung 0 = exact zero.
4449    zero_rung: bool,
4450}
4451
4452impl<'a> Q4tpView<'a> {
4453    fn new(bytes: &'a [u8], rows: usize, cols: usize) -> Self {
4454        let (params_off, codes_off, stride) = q4tp_sections(rows, cols);
4455        Self {
4456            nib: &bytes[..params_off],
4457            params: &bytes[params_off..codes_off],
4458            codes: &bytes[codes_off..],
4459            stride,
4460            zero_rung: false,
4461        }
4462    }
4463
4464    /// The q2tp view: identical params/codes planes, 8 B weight chunks.
4465    fn new_q2(bytes: &'a [u8], rows: usize, cols: usize) -> Self {
4466        let (params_off, codes_off, stride) = q2tp_sections(rows, cols);
4467        Self {
4468            nib: &bytes[..params_off],
4469            params: &bytes[params_off..codes_off],
4470            codes: &bytes[codes_off..],
4471            stride,
4472            zero_rung: true,
4473        }
4474    }
4475
4476    /// Expand row `r`'s per-tile scales into `out` (length `gpr`).
4477    ///
4478    /// Doing this once per row — rather than decoding a 5-bit code inside the
4479    /// tile loop — is what makes the format free at runtime. Random access to
4480    /// a packed 5-bit field costs a division, two bounds checks and a branch;
4481    /// the tile's actual work is two `sdot`s, so per-tile decoding dominated
4482    /// the kernel and cost 5x (measured: 1.4 vs 6.9 tok/s on Nanbeige-3B).
4483    /// Walking the plane sequentially with a bit accumulator is ~3 ops.
4484    /// Eight 5-bit codes are exactly five bytes, so a whole group of
4485    /// eight decodes from one little-endian word at fixed shifts. The
4486    /// bit-accumulator this replaces carried a data-dependent `while
4487    /// have < 5` refill whose branch sat in the innermost loop of every
4488    /// q4tp row; a decode profile put this function above the dot
4489    /// products it feeds. Same bitstream, same codes — just no branch
4490    /// and eight independent extractions.
4491    #[inline]
4492    fn scales_into(&self, r: usize, gpr: usize, out: &mut [f32]) {
4493        let tab = if self.zero_rung {
4494            q2tp_ladder(self.params, r)
4495        } else {
4496            q4tp_ladder(self.params, r)
4497        };
4498        let codes = &self.codes[r * self.stride..(r + 1) * self.stride];
4499        let out = &mut out[..gpr];
4500        let mut chunks = out.chunks_exact_mut(8);
4501        let mut ci = 0usize;
4502        for c in &mut chunks {
4503            let w = u64::from(codes[ci])
4504                | u64::from(codes[ci + 1]) << 8
4505                | u64::from(codes[ci + 2]) << 16
4506                | u64::from(codes[ci + 3]) << 24
4507                | u64::from(codes[ci + 4]) << 32;
4508            for (k, o) in c.iter_mut().enumerate() {
4509                *o = tab[((w >> (5 * k)) & 31) as usize];
4510            }
4511            ci += 5;
4512        }
4513        // Fewer than eight codes left: the shared total accessor, which
4514        // tolerates a 5-bit field whose spill byte is past the stride.
4515        let tail = &codes[ci..];
4516        for (k, o) in chunks.into_remainder().iter_mut().enumerate() {
4517            *o = tab[q4tp_code(tail, k)];
4518        }
4519    }
4520}
4521
4522#[inline]
4523fn dot_q4tp_row_i8(nib: &[u8], r: usize, gpr: usize, xq: &[i8], scales: &[f32]) -> f32 {
4524    #[cfg(target_arch = "aarch64")]
4525    unsafe {
4526        return dot_q4tp_row_sdot(nib, r, gpr, xq, scales);
4527    }
4528    #[cfg(target_arch = "x86_64")]
4529    unsafe {
4530        if vnni_tiles_enabled() {
4531            return dot_q4tp_row_vnni(nib, r, gpr, xq, scales);
4532        }
4533        return dot_q4tp_row_avx2(nib, r, gpr, xq, scales);
4534    }
4535    #[allow(unreachable_code)]
4536    {
4537        let mut acc = 0f32;
4538        for gi in 0..gpr {
4539            let tile = &nib[(r * gpr + gi) * Q4TP_NIB..(r * gpr + gi + 1) * Q4TP_NIB];
4540            let s = scales[gi];
4541            let mut d = 0i32;
4542            for (k, &b) in tile.iter().enumerate() {
4543                d += ((b & 0x0F) as i32 - 8) * xq[gi * GROUP_SIZE + k * 2] as i32
4544                    + (((b >> 4) & 0x0F) as i32 - 8) * xq[gi * GROUP_SIZE + k * 2 + 1] as i32;
4545            }
4546            acc += d as f32 * s;
4547        }
4548        acc
4549    }
4550}
4551
4552/// q4tp twin of `dot_q4t_row_sdot`: identical nibble math, but the tile
4553/// stride is 16 B (no inline scale) and the scale is a ladder lookup.
4554#[cfg(target_arch = "aarch64")]
4555#[target_feature(enable = "neon,dotprod")]
4556unsafe fn dot_q4tp_row_sdot(nib: &[u8], r: usize, gpr: usize, xq: &[i8], scales: &[f32]) -> f32 {
4557    // SAFETY: callers uphold slice-length contracts (16B tile per group,
4558    // xq.len() == gpr·GROUP_SIZE, codes covering gpr 5-bit fields).
4559    unsafe {
4560        use core::arch::aarch64::*;
4561        use core::arch::asm;
4562        let lomask = vdupq_n_u8(0x0F);
4563        let eight = vdupq_n_s8(8);
4564        let mut acc = 0f32;
4565        for gi in 0..gpr {
4566            let t = nib.as_ptr().add((r * gpr + gi) * Q4TP_NIB);
4567            let s = *scales.get_unchecked(gi);
4568            let b = vld1q_u8(t);
4569            let lo = vandq_u8(b, lomask);
4570            let hi = vshrq_n_u8::<4>(b);
4571            let e0 = vsubq_s8(vreinterpretq_s8_u8(vzip1q_u8(lo, hi)), eight);
4572            let e1 = vsubq_s8(vreinterpretq_s8_u8(vzip2q_u8(lo, hi)), eight);
4573            let x0 = vld1q_s8(xq.as_ptr().add(gi * GROUP_SIZE));
4574            let x1 = vld1q_s8(xq.as_ptr().add(gi * GROUP_SIZE + 16));
4575            let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
4576            asm!(
4577                "sdot {a0:v}.4s, {e0:v}.16b, {x0:v}.16b",
4578                "sdot {a1:v}.4s, {e1:v}.16b, {x1:v}.16b",
4579                a0 = inout(vreg) a0, a1 = inout(vreg) a1,
4580                e0 = in(vreg) e0, x0 = in(vreg) x0, e1 = in(vreg) e1, x1 = in(vreg) x1,
4581                options(pure, nomem, nostack),
4582            );
4583            acc += vaddvq_s32(vaddq_s32(a0, a1)) as f32 * s;
4584        }
4585        acc
4586    }
4587}
4588
4589#[cfg(target_arch = "x86_64")]
4590#[target_feature(enable = "avx2")]
4591unsafe fn dot_q4tp_row_avx2(nib: &[u8], r: usize, gpr: usize, xq: &[i8], scales: &[f32]) -> f32 {
4592    // SAFETY: see dot_q4tp_row_sdot.
4593    unsafe {
4594        use core::arch::x86_64::*;
4595        let lomask = _mm_set1_epi8(0x0F);
4596        let eight = _mm256_set1_epi8(8);
4597        let ones = _mm256_set1_epi16(1);
4598        let mut acc = 0f32;
4599        for gi in 0..gpr {
4600            let t = nib.as_ptr().add((r * gpr + gi) * Q4TP_NIB);
4601            let s = *scales.get_unchecked(gi);
4602            let b = _mm_loadu_si128(t as *const __m128i);
4603            let lo = _mm_and_si128(b, lomask);
4604            let hi = _mm_and_si128(_mm_srli_epi16::<4>(b), lomask);
4605            let w = _mm256_sub_epi8(
4606                _mm256_set_m128i(_mm_unpackhi_epi8(lo, hi), _mm_unpacklo_epi8(lo, hi)),
4607                eight,
4608            );
4609            let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
4610            let p16 = _mm256_maddubs_epi16(_mm256_abs_epi8(w), _mm256_sign_epi8(x, w));
4611            let d = _mm256_madd_epi16(p16, ones);
4612            let hi128 = _mm256_extracti128_si256::<1>(d);
4613            let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
4614            let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
4615            let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
4616            acc += _mm_cvtsi128_si32(s32) as f32 * s;
4617        }
4618        acc
4619    }
4620}
4621
4622/// VNNI twin of `dot_q4tp_row_avx2` (see `dot_q4t_row_vnni` for why the
4623/// 256-bit VL encoding is the one to use here).
4624#[cfg(target_arch = "x86_64")]
4625#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
4626unsafe fn dot_q4tp_row_vnni(nib: &[u8], r: usize, gpr: usize, xq: &[i8], scales: &[f32]) -> f32 {
4627    // SAFETY: see dot_q4tp_row_sdot.
4628    unsafe {
4629        use core::arch::x86_64::*;
4630        let lomask = _mm_set1_epi8(0x0F);
4631        let eight = _mm256_set1_epi8(8);
4632        let mut acc = 0f32;
4633        for gi in 0..gpr {
4634            let t = nib.as_ptr().add((r * gpr + gi) * Q4TP_NIB);
4635            let s = *scales.get_unchecked(gi);
4636            let b = _mm_loadu_si128(t as *const __m128i);
4637            let lo = _mm_and_si128(b, lomask);
4638            let hi = _mm_and_si128(_mm_srli_epi16::<4>(b), lomask);
4639            let w = _mm256_sub_epi8(
4640                _mm256_set_m128i(_mm_unpackhi_epi8(lo, hi), _mm_unpacklo_epi8(lo, hi)),
4641                eight,
4642            );
4643            let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
4644            acc += dpbusd_hsum(_mm256_abs_epi8(w), _mm256_sign_epi8(x, w)) as f32 * s;
4645        }
4646        acc
4647    }
4648}
4649
4650/// Exact scalar q4tp row — the `CMF_SDOT=0` contract, same pairwise
4651/// accumulation shape as `q4t_row_exact`.
4652#[inline]
4653fn q4tp_row_exact(nib: &[u8], r: usize, gpr: usize, x: &[f32], scales: &[f32]) -> f32 {
4654    let mut acc = 0f32;
4655    for gi in 0..gpr {
4656        let tile = &nib[(r * gpr + gi) * Q4TP_NIB..(r * gpr + gi + 1) * Q4TP_NIB];
4657        let s = scales[gi];
4658        let xg = &x[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
4659        let mut ga = 0f32;
4660        for (k, &b) in tile.iter().enumerate() {
4661            ga += ((b & 0x0F) as f32 - 8.0) * xg[k * 2]
4662                + (((b >> 4) & 0x0F) as f32 - 8.0) * xg[k * 2 + 1];
4663        }
4664        acc += ga * s;
4665    }
4666    acc
4667}
4668
4669/// Single weight of a q4tp tensor — the a8w8 outlier path, which restores
4670/// activation outliers at full precision after the int8 pass.
4671#[inline]
4672fn q4tp_outlier(nib: &[u8], r: usize, gpr: usize, j: usize, scales: &[f32]) -> (f32, f32) {
4673    let (gi, k) = (j / GROUP_SIZE, j % GROUP_SIZE);
4674    let byte = nib[(r * gpr + gi) * Q4TP_NIB + k / 2];
4675    let n = if k & 1 == 0 { byte & 0x0F } else { byte >> 4 };
4676    ((n as i32 - 8) as f32, scales[gi])
4677}
4678
4679/// Fused q4tp matvec (dispatch mirrors `q4t_matvec`).
4680fn q4tp_matvec(
4681    bytes: &[u8],
4682    x: &[f32],
4683    rows: usize,
4684    cols: usize,
4685    out: &mut [f32],
4686    pool: Option<&Pool>,
4687) {
4688    debug_assert_eq!(out.len(), rows);
4689    let gpr = cols / GROUP_SIZE;
4690    let v = Q4tpView::new(bytes, rows, cols);
4691    let out_addr = SendMut(out.as_mut_ptr());
4692    if a8w8_enabled() {
4693        let act = split_act(x);
4694        let run = |start: usize, end: usize| {
4695            // One scratch row of scales per worker — borrowed, not minted.
4696            with_krow(gpr, |sc| {
4697                for r in start..end {
4698                    v.scales_into(r, gpr, sc);
4699                    let mut acc = dot_q4tp_row_i8(v.nib, r, gpr, &act.xq, sc) * act.sx;
4700                    for &(j, xv) in &act.outliers {
4701                        let (w, s) = q4tp_outlier(v.nib, r, gpr, j, sc);
4702                        acc += w * s * xv;
4703                    }
4704                    // SAFETY: disjoint row ranges per worker.
4705                    unsafe { *out_addr.at(r) = acc };
4706                }
4707            })
4708        };
4709        dispatch_rows(pool, rows, &run);
4710        return;
4711    }
4712    let run = |start: usize, end: usize| {
4713        with_krow(gpr, |sc| {
4714            for r in start..end {
4715                v.scales_into(r, gpr, sc);
4716                // SAFETY: disjoint row ranges per worker.
4717                unsafe { *out_addr.at(r) = q4tp_row_exact(v.nib, r, gpr, x, sc) };
4718            }
4719        })
4720    };
4721    dispatch_rows(pool, rows, &run);
4722}
4723
4724/// Fused two-input q4tp matvec — the SwiGLU gate/up pair. Weights and the
4725/// row ladder are read once and spent on both activation streams.
4726#[allow(clippy::too_many_arguments)]
4727fn q4tp_matvec2(
4728    bytes: &[u8],
4729    x1: &[f32],
4730    x2: &[f32],
4731    rows: usize,
4732    cols: usize,
4733    o1: &mut [f32],
4734    o2: &mut [f32],
4735    pool: Option<&Pool>,
4736) {
4737    let gpr = cols / GROUP_SIZE;
4738    let v = Q4tpView::new(bytes, rows, cols);
4739    let (p1, p2) = (SendMut(o1.as_mut_ptr()), SendMut(o2.as_mut_ptr()));
4740    let run = |start: usize, end: usize| {
4741        let mut sc = vec![0f32; gpr];
4742        for r in start..end {
4743            v.scales_into(r, gpr, &mut sc);
4744            // SAFETY: disjoint row ranges per worker.
4745            unsafe {
4746                *p1.at(r) = q4tp_row_exact(v.nib, r, gpr, x1, &sc);
4747                *p2.at(r) = q4tp_row_exact(v.nib, r, gpr, x2, &sc);
4748            }
4749        }
4750    };
4751    dispatch_rows(pool, rows, &run);
4752}
4753
4754/// One q2tp outlier weight at column `j` of row `r`: the 2-bit code and
4755/// its group scale, mirrored on `q4tp_outlier`.
4756#[inline]
4757fn q2tp_outlier(chunks: &[u8], r: usize, gpr: usize, j: usize, scales: &[f32]) -> (f32, f32) {
4758    let (gi, k) = (j / GROUP_SIZE, j % GROUP_SIZE);
4759    let byte = chunks[(r * gpr + gi) * Q2TP_CHUNK + k / 4];
4760    let c = (byte >> (2 * (k % 4))) & 3;
4761    (c as f32 - 1.5, scales[gi])
4762}
4763
4764#[cfg(target_arch = "x86_64")]
4765const Q2TP_DECODE_U32: [u32; 256] = {
4766    let mut tab = [0u32; 256];
4767    let mut b = 0usize;
4768    while b < 256 {
4769        tab[b] = ((b as u32) & 3)
4770            | ((((b as u32) >> 2) & 3) << 8)
4771            | ((((b as u32) >> 4) & 3) << 16)
4772            | ((((b as u32) >> 6) & 3) << 24);
4773        b += 1;
4774    }
4775    tab
4776};
4777
4778/// Eight packed q2tp bytes against 32 signed activation bytes. `maddubs`
4779/// exactly computes unsigned 2-bit code × signed i8; its pair sums cannot
4780/// saturate (2 × 3 × 127 < i16::MAX), and the second madd widens to i32.
4781#[cfg(target_arch = "x86_64")]
4782#[target_feature(enable = "avx2")]
4783unsafe fn q2tp_code_dot_avx2(ch: &[u8], x: &[i8]) -> i32 {
4784    use core::arch::x86_64::*;
4785    debug_assert!(ch.len() >= Q2TP_CHUNK && x.len() >= GROUP_SIZE);
4786    let codes = _mm256_setr_epi32(
4787        Q2TP_DECODE_U32[ch[0] as usize] as i32,
4788        Q2TP_DECODE_U32[ch[1] as usize] as i32,
4789        Q2TP_DECODE_U32[ch[2] as usize] as i32,
4790        Q2TP_DECODE_U32[ch[3] as usize] as i32,
4791        Q2TP_DECODE_U32[ch[4] as usize] as i32,
4792        Q2TP_DECODE_U32[ch[5] as usize] as i32,
4793        Q2TP_DECODE_U32[ch[6] as usize] as i32,
4794        Q2TP_DECODE_U32[ch[7] as usize] as i32,
4795    );
4796    let xv = unsafe { _mm256_loadu_si256(x.as_ptr().cast()) };
4797    let pair = _mm256_maddubs_epi16(codes, xv);
4798    let quad = _mm256_madd_epi16(pair, _mm256_set1_epi16(1));
4799    let sum128 = _mm_add_epi32(
4800        _mm256_castsi256_si128(quad),
4801        _mm256_extracti128_si256(quad, 1),
4802    );
4803    let sum64 = _mm_hadd_epi32(sum128, sum128);
4804    _mm_cvtsi128_si32(_mm_hadd_epi32(sum64, sum64))
4805}
4806
4807/// Integer dot of one q2tp row against pre-quantized activations:
4808/// Σ_g s_g · (Σ c·xq − 1.5·Σ xq). The half-integer grid (c − 1.5)
4809/// becomes exact integer math through the group sums — the same trick
4810/// every a8w8 kernel in this file rides. The codes decode into a
4811/// 32-byte scratch in natural order and the dot itself is the shared
4812/// SDOT primitive; elsewhere a scalar integer loop.
4813#[inline]
4814fn dot_q2tp_row_i8(
4815    chunks: &[u8],
4816    r: usize,
4817    gpr: usize,
4818    xq: &[i8],
4819    gsum: &[i32],
4820    scales: &[f32],
4821) -> f32 {
4822    let mut acc = 0f32;
4823    let base = r * gpr * Q2TP_CHUNK;
4824    #[cfg(not(any(target_arch = "aarch64", target_arch = "x86_64")))]
4825    let mut codes = [0i8; GROUP_SIZE];
4826    #[cfg(target_arch = "x86_64")]
4827    let avx2 = std::arch::is_x86_feature_detected!("avx2");
4828    for gi in 0..gpr {
4829        let ch = &chunks[base + gi * Q2TP_CHUNK..base + (gi + 1) * Q2TP_CHUNK];
4830        let xg = &xq[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
4831        #[cfg(target_arch = "aarch64")]
4832        // NEON: the byte's four 2-bit fields land in four lane vectors
4833        // (shift+mask), vld4 de-interleaves xq to match (xj[k] =
4834        // xq[4k+j]), widening MACs accumulate exactly in i32. A scalar
4835        // decode here cost as much as the dot it fed — the profile put
4836        // it at the top of the whole W2 decode.
4837        let dot = unsafe {
4838            use core::arch::aarch64::*;
4839            let b = vld1_u8(ch.as_ptr());
4840            let three = vdup_n_u8(3);
4841            let c0 = vreinterpret_s8_u8(vand_u8(b, three));
4842            let c1 = vreinterpret_s8_u8(vand_u8(vshr_n_u8(b, 2), three));
4843            let c2 = vreinterpret_s8_u8(vand_u8(vshr_n_u8(b, 4), three));
4844            let c3 = vreinterpret_s8_u8(vand_u8(vshr_n_u8(b, 6), three));
4845            let x4 = vld4_s8(xg.as_ptr());
4846            let mut acc4 = vdupq_n_s32(0);
4847            acc4 = vpadalq_s16(acc4, vmull_s8(c0, x4.0));
4848            acc4 = vpadalq_s16(acc4, vmull_s8(c1, x4.1));
4849            acc4 = vpadalq_s16(acc4, vmull_s8(c2, x4.2));
4850            acc4 = vpadalq_s16(acc4, vmull_s8(c3, x4.3));
4851            vaddvq_s32(acc4)
4852        };
4853        #[cfg(target_arch = "x86_64")]
4854        let dot: i32 = if avx2 {
4855            // SAFETY: the runtime feature check gates the target-feature body;
4856            // the group slices above are exactly 8 and 32 bytes long.
4857            unsafe { q2tp_code_dot_avx2(ch, xg) }
4858        } else {
4859            ch.iter()
4860                .enumerate()
4861                .map(|(k, &b)| {
4862                    ((b & 3) as i32) * xg[k * 4] as i32
4863                        + (((b >> 2) & 3) as i32) * xg[k * 4 + 1] as i32
4864                        + (((b >> 4) & 3) as i32) * xg[k * 4 + 2] as i32
4865                        + (((b >> 6) & 3) as i32) * xg[k * 4 + 3] as i32
4866                })
4867                .sum()
4868        };
4869        #[cfg(not(any(target_arch = "aarch64", target_arch = "x86_64")))]
4870        let dot: i32 = {
4871            for (k, &b) in ch.iter().enumerate() {
4872                codes[k * 4] = (b & 3) as i8;
4873                codes[k * 4 + 1] = ((b >> 2) & 3) as i8;
4874                codes[k * 4 + 2] = ((b >> 4) & 3) as i8;
4875                codes[k * 4 + 3] = ((b >> 6) & 3) as i8;
4876            }
4877            codes
4878                .iter()
4879                .zip(xg)
4880                .map(|(&c, &x)| c as i32 * x as i32)
4881                .sum()
4882        };
4883        acc += scales[gi] * (dot as f32 - 1.5 * gsum[gi] as f32);
4884    }
4885    acc
4886}
4887
4888/// Exact f32 dot of one q2tp row: 2-bit fields LSB-first, (c − 1.5)·s.
4889/// Scalar on purpose — the 2-bit class targets the GPU graph; the CPU
4890/// path exists for parity gates and small-machine fallback.
4891fn q2tp_row_exact(chunks: &[u8], r: usize, gpr: usize, x: &[f32], scales: &[f32]) -> f32 {
4892    let mut acc = 0f32;
4893    for gi in 0..gpr {
4894        let ch = &chunks[(r * gpr + gi) * Q2TP_CHUNK..(r * gpr + gi + 1) * Q2TP_CHUNK];
4895        let s = scales[gi];
4896        let xb = &x[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
4897        let mut g = 0f32;
4898        for (k, &b) in ch.iter().enumerate() {
4899            g += ((b & 3) as f32 - 1.5) * xb[k * 4]
4900                + (((b >> 2) & 3) as f32 - 1.5) * xb[k * 4 + 1]
4901                + (((b >> 4) & 3) as f32 - 1.5) * xb[k * 4 + 2]
4902                + (((b >> 6) & 3) as f32 - 1.5) * xb[k * 4 + 3];
4903        }
4904        acc += s * g;
4905    }
4906    acc
4907}
4908
4909fn q2tp_matvec(
4910    bytes: &[u8],
4911    x: &[f32],
4912    rows: usize,
4913    cols: usize,
4914    out: &mut [f32],
4915    pool: Option<&Pool>,
4916) {
4917    debug_assert_eq!(out.len(), rows);
4918    let gpr = cols / GROUP_SIZE;
4919    let v = Q4tpView::new_q2(bytes, rows, cols);
4920    let out_addr = SendMut(out.as_mut_ptr());
4921    // a8w8 fast path (CMF_SDOT=0 keeps the exact scalar walk): integer
4922    // code dots + group sums, exact outlier correction — the same
4923    // contract as every sibling kernel; measured 2-bit rows were the
4924    // only scalar holdout in the family.
4925    if a8w8_enabled() {
4926        let act = split_act(x);
4927        let gsum = q1_group_sums(&act.xq, gpr);
4928        let (act, gsum) = (&act, &gsum);
4929        let run = move |start: usize, end: usize| {
4930            with_krow(gpr, |sc| {
4931                for r in start..end {
4932                    v.scales_into(r, gpr, sc);
4933                    let mut acc = dot_q2tp_row_i8(v.nib, r, gpr, &act.xq, gsum, sc) * act.sx;
4934                    for &(j, xv) in &act.outliers {
4935                        let (w, s) = q2tp_outlier(v.nib, r, gpr, j, sc);
4936                        acc += w * s * xv;
4937                    }
4938                    // SAFETY: disjoint row ranges per worker.
4939                    unsafe { *out_addr.at(r) = acc };
4940                }
4941            })
4942        };
4943        dispatch_rows(pool, rows, &run);
4944        return;
4945    }
4946    let run = |start: usize, end: usize| {
4947        with_krow(gpr, |sc| {
4948            for r in start..end {
4949                v.scales_into(r, gpr, sc);
4950                // SAFETY: disjoint row ranges per worker.
4951                unsafe { *out_addr.at(r) = q2tp_row_exact(v.nib, r, gpr, x, sc) };
4952            }
4953        })
4954    };
4955    dispatch_rows(pool, rows, &run);
4956}
4957
4958/// Fused two-input q2tp matvec — the SwiGLU gate/up pair.
4959#[allow(clippy::too_many_arguments)]
4960fn q2tp_matvec2(
4961    bytes: &[u8],
4962    x1: &[f32],
4963    x2: &[f32],
4964    rows: usize,
4965    cols: usize,
4966    o1: &mut [f32],
4967    o2: &mut [f32],
4968    pool: Option<&Pool>,
4969) {
4970    let gpr = cols / GROUP_SIZE;
4971    let v = Q4tpView::new_q2(bytes, rows, cols);
4972    let (p1, p2) = (SendMut(o1.as_mut_ptr()), SendMut(o2.as_mut_ptr()));
4973    let run = |start: usize, end: usize| {
4974        let mut sc = vec![0f32; gpr];
4975        for r in start..end {
4976            v.scales_into(r, gpr, &mut sc);
4977            // SAFETY: disjoint row ranges per worker.
4978            unsafe {
4979                *p1.at(r) = q2tp_row_exact(v.nib, r, gpr, x1, &sc);
4980                *p2.at(r) = q2tp_row_exact(v.nib, r, gpr, x2, &sc);
4981            }
4982        }
4983    };
4984    dispatch_rows(pool, rows, &run);
4985}
4986
4987/// Batched q2tp matmat: scalar row kernel over every batch column. CPU
4988/// prefill only — decode rides the graph, so plain and correct beats
4989/// clever here.
4990/// Test doors into the host 2-bit kernels: the stand's heap corruption
4991/// pointed at down-shaped tensors, and the private fns need a way to be
4992/// held to a reference without a model file around them.
4993pub fn q2tp_matvec_for_test(bytes: &[u8], x: &[f32], rows: usize, cols: usize, out: &mut [f32]) {
4994    // The facade IS the reference: encoder oracles hold requant output
4995    // to the exact scalar walk. The production dispatch may take the i8
4996    // fast path, whose error scale is the ACTIVATIONS' — a different
4997    // claim than the encoder correctness these tests pin.
4998    let gpr = cols / GROUP_SIZE;
4999    let v = Q4tpView::new_q2(bytes, rows, cols);
5000    with_krow(gpr, |sc| {
5001        for r in 0..rows {
5002            v.scales_into(r, gpr, sc);
5003            out[r] = q2tp_row_exact(v.nib, r, gpr, x, sc);
5004        }
5005    });
5006}
5007
5008pub fn q2tp_matmat_for_test(
5009    bytes: &[u8],
5010    xs_all: &[f32],
5011    b: usize,
5012    rows: usize,
5013    cols: usize,
5014    out: &mut [f32],
5015) {
5016    q2tp_matmat(bytes, xs_all, b, rows, cols, out, None);
5017}
5018
5019fn q2tp_matmat(
5020    bytes: &[u8],
5021    xs_all: &[f32],
5022    b: usize,
5023    rows: usize,
5024    cols: usize,
5025    out: &mut [f32],
5026    pool: Option<&Pool>,
5027) {
5028    debug_assert_eq!(out.len(), b * rows);
5029    let gpr = cols / GROUP_SIZE;
5030    let v = Q4tpView::new_q2(bytes, rows, cols);
5031    let out_addr = SendMut(out.as_mut_ptr());
5032    let run = |start: usize, end: usize| {
5033        let mut sc = vec![0f32; gpr];
5034        for r in start..end {
5035            v.scales_into(r, gpr, &mut sc);
5036            for bi in 0..b {
5037                let x = &xs_all[bi * cols..(bi + 1) * cols];
5038                // SAFETY: disjoint row ranges per worker.
5039                unsafe { *out_addr.at(bi * rows + r) = q2tp_row_exact(v.nib, r, gpr, x, &sc) };
5040            }
5041        }
5042    };
5043    dispatch_rows(pool, rows, &run);
5044}
5045
5046/// The pre-vectorised shape, kept for A/B (`CMF_Q4TP_V1=1`): the
5047/// horizontal add lands once per group per column instead of once per
5048/// row. Same weights, same activations — only the reduction differs.
5049#[cfg(target_arch = "aarch64")]
5050#[target_feature(enable = "neon,dotprod")]
5051unsafe fn dot_q4tp_row_1x4_sdot_v1(
5052    nib: &[u8],
5053    r: usize,
5054    gpr: usize,
5055    xs: [&[i8]; 4],
5056    scales: &[f32],
5057) -> [f32; 4] {
5058    unsafe {
5059        use core::arch::aarch64::*;
5060        use core::arch::asm;
5061        let lomask = vdupq_n_u8(0x0F);
5062        let eight = vdupq_n_s8(8);
5063        let (mut f0, mut f1, mut f2, mut f3) = (0f32, 0f32, 0f32, 0f32);
5064        for gi in 0..gpr {
5065            let t = nib.as_ptr().add((r * gpr + gi) * Q4TP_NIB);
5066            let s = *scales.get_unchecked(gi);
5067            let bb = vld1q_u8(t);
5068            let lo = vandq_u8(bb, lomask);
5069            let hi = vshrq_n_u8::<4>(bb);
5070            let e0 = vsubq_s8(vreinterpretq_s8_u8(vzip1q_u8(lo, hi)), eight);
5071            let e1 = vsubq_s8(vreinterpretq_s8_u8(vzip2q_u8(lo, hi)), eight);
5072            let mut d = [0f32; 4];
5073            for (k, dk) in d.iter_mut().enumerate() {
5074                let x0 = vld1q_s8(xs[k].as_ptr().add(gi * GROUP_SIZE));
5075                let x1 = vld1q_s8(xs[k].as_ptr().add(gi * GROUP_SIZE + 16));
5076                let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
5077                asm!(
5078                    "sdot {a0:v}.4s, {e0:v}.16b, {x0:v}.16b",
5079                    "sdot {a1:v}.4s, {e1:v}.16b, {x1:v}.16b",
5080                    a0 = inout(vreg) a0, a1 = inout(vreg) a1,
5081                    e0 = in(vreg) e0, x0 = in(vreg) x0, e1 = in(vreg) e1, x1 = in(vreg) x1,
5082                    options(pure, nomem, nostack),
5083                );
5084                *dk = vaddvq_s32(vaddq_s32(a0, a1)) as f32 * s;
5085            }
5086            f0 += d[0];
5087            f1 += d[1];
5088            f2 += d[2];
5089            f3 += d[3];
5090        }
5091        [f0, f1, f2, f3]
5092    }
5093}
5094
5095/// Which q4tp batch kernel to run: 1 = the previous one, 2 = the tuned
5096/// one, 0 = decide from the CPU. An atomic rather than a `OnceLock` so a
5097/// benchmark can alternate the two inside one process, where the machine's
5098/// mood — a shared box drifts ±25% between runs — is the same for both.
5099/// What the two mean is per-architecture: on x86 the blocked AVX-512 path
5100/// against the per-column one, on ARM the two reduction shapes.
5101#[allow(dead_code)]
5102static Q4TP_ALT: std::sync::atomic::AtomicU8 = std::sync::atomic::AtomicU8::new(0);
5103
5104/// Blocking pays on x86 only with 512-bit VNNI. With AVX2 alone, four
5105/// columns sharing an unpack still measured slower than the per-column
5106/// path (23.2 ms against 19.4 on a 48-thread EPYC), because that path
5107/// already dequantizes the row once — so the blocked kernel bought a
5108/// second unpack-free pass at the price of half the vector width.
5109#[cfg(target_arch = "x86_64")]
5110fn q4tp_blocked_x86() -> bool {
5111    match Q4TP_ALT.load(std::sync::atomic::Ordering::Relaxed) {
5112        1 => false,
5113        // A forced ON still asks the CPU. The switch exists so a bench can
5114        // pick a kernel, not so it can promise instructions the machine
5115        // does not have — CI caught that as a SIGILL on a runner without
5116        // AVX-512, where the parity test had turned the path on by hand.
5117        2 => avx512vnni_enabled(),
5118        // Deliberately not cached back into the switch: both gates below
5119        // hold their own `OnceLock`, and latching their answer here would
5120        // make a test's override outlive the test that set it.
5121        _ => blocked_enabled() && avx512vnni_enabled(),
5122    }
5123}
5124
5125/// `CMF_Q4TP_V1=1` picks the old reduction shape (A/B only).
5126#[cfg(target_arch = "aarch64")]
5127#[allow(dead_code)]
5128fn q4tp_v1() -> bool {
5129    match Q4TP_ALT.load(std::sync::atomic::Ordering::Relaxed) {
5130        1 => true,
5131        2 => false,
5132        _ => {
5133            static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
5134            *ON.get_or_init(|| std::env::var("CMF_Q4TP_V1").is_ok_and(|v| v != "0"))
5135        }
5136    }
5137}
5138
5139/// Two weight rows against eight columns. The activation load is the
5140/// same for both rows, so it is paid once for twice the arithmetic, and
5141/// sixteen accumulator chains run where eight did — which is what a kernel
5142/// retiring 0.29 instructions a cycle is short of. Register pressure is
5143/// the limit: sixteen `zmm` accumulators, two weight tiles, one
5144/// activation, of thirty-two.
5145///
5146/// Four rows by four columns spends the same sixteen accumulators the
5147/// other way and measured worse — 1488 GFLOP/s against 1644 — so the
5148/// unpack, which four rows pay twice as often, costs more than the extra
5149/// sharing of one activation load buys.
5150#[cfg(target_arch = "x86_64")]
5151#[target_feature(enable = "avx512f,avx512bw,avx512vnni")]
5152unsafe fn dot_q4tp_2x8_avx512(
5153    nib: &[u8],
5154    r0: usize,
5155    gpr: usize,
5156    xs: [&[i8]; 8],
5157    sc0: &[f32],
5158    sc1: &[f32],
5159) -> [[f32; 8]; 2] {
5160    // SAFETY: as dot_q4tp_row_1x8_avx512, two adjacent rows at once; the
5161    // caller guarantees r0 + 1 < rows and the ISA.
5162    unsafe {
5163        use core::arch::x86_64::*;
5164        let lomask = _mm256_set1_epi8(0x0F);
5165        let eight = _mm256_set1_epi8(8);
5166        let zero = _mm512_setzero_si512();
5167        let mut v0 = [_mm512_setzero_ps(); 8];
5168        let mut v1 = [_mm512_setzero_ps(); 8];
5169        let pairs = gpr / 2;
5170        let unpack = |r: usize, gi: usize| -> (__m512i, __mmask64) {
5171            let t = nib.as_ptr().add((r * gpr + gi) * Q4TP_NIB);
5172            let bb = _mm256_loadu_si256(t as *const __m256i);
5173            let lo = _mm256_and_si256(bb, lomask);
5174            let hi = _mm256_and_si256(_mm256_srli_epi16::<4>(bb), lomask);
5175            let ul = _mm256_sub_epi8(_mm256_unpacklo_epi8(lo, hi), eight);
5176            let uh = _mm256_sub_epi8(_mm256_unpackhi_epi8(lo, hi), eight);
5177            let cat = _mm512_inserti64x4::<1>(_mm512_castsi256_si512(ul), uh);
5178            let w = _mm512_shuffle_i64x2::<0b11_01_10_00>(cat, cat);
5179            (_mm512_abs_epi8(w), _mm512_movepi8_mask(w))
5180        };
5181        for gp in 0..pairs {
5182            let gi = gp * 2;
5183            let (wa0, neg0) = unpack(r0, gi);
5184            let (wa1, neg1) = unpack(r0 + 1, gi);
5185            let off = gi * GROUP_SIZE;
5186            let sv = |sc: &[f32]| {
5187                _mm512_insertf32x8::<1>(
5188                    _mm512_castps256_ps512(_mm256_set1_ps(*sc.get_unchecked(gi))),
5189                    _mm256_set1_ps(*sc.get_unchecked(gi + 1)),
5190                )
5191            };
5192            let s0 = sv(sc0);
5193            let s1 = sv(sc1);
5194            for k in 0..8 {
5195                let xv = _mm512_loadu_si512(xs[k].as_ptr().add(off) as *const __m512i);
5196                let d0 = _mm512_cvtepi32_ps(_mm512_dpbusd_epi32(
5197                    zero,
5198                    wa0,
5199                    _mm512_mask_sub_epi8(xv, neg0, zero, xv),
5200                ));
5201                let d1 = _mm512_cvtepi32_ps(_mm512_dpbusd_epi32(
5202                    zero,
5203                    wa1,
5204                    _mm512_mask_sub_epi8(xv, neg1, zero, xv),
5205                ));
5206                v0[k] = _mm512_fmadd_ps(d0, s0, v0[k]);
5207                v1[k] = _mm512_fmadd_ps(d1, s1, v1[k]);
5208            }
5209        }
5210        let mut acc = [[0f32; 8]; 2];
5211        for k in 0..8 {
5212            acc[0][k] = _mm512_reduce_add_ps(v0[k]);
5213            acc[1][k] = _mm512_reduce_add_ps(v1[k]);
5214        }
5215        if gpr % 2 == 1 {
5216            let off = (gpr - 1) * GROUP_SIZE;
5217            for j in off..off + GROUP_SIZE {
5218                let (w0, sa) = q4tp_outlier(nib, r0, gpr, j, sc0);
5219                let (w1, sb) = q4tp_outlier(nib, r0 + 1, gpr, j, sc1);
5220                for k in 0..8 {
5221                    let x = *xs[k].get_unchecked(j) as f32;
5222                    acc[0][k] += w0 * sa * x;
5223                    acc[1][k] += w1 * sb * x;
5224                }
5225            }
5226        }
5227        acc
5228    }
5229}
5230
5231/// The same, eight columns at a time. One unpack then feeds twice as many
5232/// activation streams, so a wide batch reads the weight tile half as
5233/// often; the price is eight accumulators live at once. Measured 9.0 ->
5234/// 8.3 ms at 9216x2304, b=296 on a 48-thread EPYC 9B45.
5235#[cfg(target_arch = "x86_64")]
5236#[target_feature(enable = "avx512f,avx512bw,avx512vnni")]
5237unsafe fn dot_q4tp_row_1x8_avx512(
5238    nib: &[u8],
5239    r: usize,
5240    gpr: usize,
5241    xs: [&[i8]; 8],
5242    scales: &[f32],
5243) -> [f32; 8] {
5244    // SAFETY: as dot_q4tp_row_1x4_avx2; caller guarantees the ISA.
5245    unsafe {
5246        use core::arch::x86_64::*;
5247        let lomask = _mm256_set1_epi8(0x0F);
5248        let eight = _mm256_set1_epi8(8);
5249        let zero = _mm512_setzero_si512();
5250        let (mut v0, mut v1, mut v2, mut v3) = (
5251            _mm512_setzero_ps(),
5252            _mm512_setzero_ps(),
5253            _mm512_setzero_ps(),
5254            _mm512_setzero_ps(),
5255        );
5256        let (mut v4, mut v5, mut v6, mut v7) = (
5257            _mm512_setzero_ps(),
5258            _mm512_setzero_ps(),
5259            _mm512_setzero_ps(),
5260            _mm512_setzero_ps(),
5261        );
5262        let pairs = gpr / 2;
5263        for gp in 0..pairs {
5264            let gi = gp * 2;
5265            let t = nib.as_ptr().add((r * gpr + gi) * Q4TP_NIB);
5266            let bb = _mm256_loadu_si256(t as *const __m256i);
5267            let lo = _mm256_and_si256(bb, lomask);
5268            let hi = _mm256_and_si256(_mm256_srli_epi16::<4>(bb), lomask);
5269            // `unpack` works per 128-bit lane, so the halves come out as
5270            // [A.lo, B.lo] and [A.hi, B.hi]; the shuffle reorders the four
5271            // 128-bit lanes into the weights' natural order, which is what
5272            // the straight activation load expects.
5273            let ul = _mm256_sub_epi8(_mm256_unpacklo_epi8(lo, hi), eight);
5274            let uh = _mm256_sub_epi8(_mm256_unpackhi_epi8(lo, hi), eight);
5275            let cat = _mm512_inserti64x4::<1>(_mm512_castsi256_si512(ul), uh);
5276            let w = _mm512_shuffle_i64x2::<0b11_01_10_00>(cat, cat);
5277            let wabs = _mm512_abs_epi8(w);
5278            let neg = _mm512_movepi8_mask(w);
5279            let off = gi * GROUP_SIZE;
5280            let sv = _mm512_insertf32x8::<1>(
5281                _mm512_castps256_ps512(_mm256_set1_ps(*scales.get_unchecked(gi))),
5282                _mm256_set1_ps(*scales.get_unchecked(gi + 1)),
5283            );
5284            let dot = |x: &[i8]| -> __m512 {
5285                let xv = _mm512_loadu_si512(x.as_ptr().add(off) as *const __m512i);
5286                let sx = _mm512_mask_sub_epi8(xv, neg, zero, xv);
5287                _mm512_cvtepi32_ps(_mm512_dpbusd_epi32(zero, wabs, sx))
5288            };
5289            v0 = _mm512_fmadd_ps(dot(xs[0]), sv, v0);
5290            v1 = _mm512_fmadd_ps(dot(xs[1]), sv, v1);
5291            v2 = _mm512_fmadd_ps(dot(xs[2]), sv, v2);
5292            v3 = _mm512_fmadd_ps(dot(xs[3]), sv, v3);
5293            v4 = _mm512_fmadd_ps(dot(xs[4]), sv, v4);
5294            v5 = _mm512_fmadd_ps(dot(xs[5]), sv, v5);
5295            v6 = _mm512_fmadd_ps(dot(xs[6]), sv, v6);
5296            v7 = _mm512_fmadd_ps(dot(xs[7]), sv, v7);
5297        }
5298        let mut acc = [
5299            _mm512_reduce_add_ps(v0),
5300            _mm512_reduce_add_ps(v1),
5301            _mm512_reduce_add_ps(v2),
5302            _mm512_reduce_add_ps(v3),
5303            _mm512_reduce_add_ps(v4),
5304            _mm512_reduce_add_ps(v5),
5305            _mm512_reduce_add_ps(v6),
5306            _mm512_reduce_add_ps(v7),
5307        ];
5308        // An odd group count leaves one group over; the narrow kernel
5309        // finishes it rather than the tail being a special case here.
5310        if gpr % 2 == 1 {
5311            let off = (gpr - 1) * GROUP_SIZE;
5312            for j in off..off + GROUP_SIZE {
5313                let (w, s) = q4tp_outlier(nib, r, gpr, j, scales);
5314                let ws = w * s;
5315                for k in 0..8 {
5316                    acc[k] += ws * *xs[k].get_unchecked(j) as f32;
5317                }
5318            }
5319        }
5320        acc
5321    }
5322}
5323
5324/// The same four columns, 512 bits wide. Two groups (64 weights) ride one
5325/// unpack and one `vpdpbusd`, where AVX2 needs two unpacks and four
5326/// `maddubs`/`madd` pairs — about 2.3x fewer instructions for the same
5327/// arithmetic. The two groups carry different scales, so the fma takes a
5328/// vector whose halves hold each group's scale rather than a broadcast.
5329///
5330/// There is no 512-bit `vpsignb`, so the activation's sign is applied by
5331/// negating under a mask taken from the weight's sign bits. That mask is
5332/// per-tile, so it is hoisted out of the column loop and the per-column
5333/// cost stays exactly one instruction, as with `sign_epi8`. Weights of
5334/// zero are not zeroed by the mask trick and do not need to be: their
5335/// magnitude is zero, so the product is.
5336#[cfg(target_arch = "x86_64")]
5337#[target_feature(enable = "avx512f,avx512bw,avx512vnni")]
5338unsafe fn dot_q4tp_row_1x4_avx512(
5339    nib: &[u8],
5340    r: usize,
5341    gpr: usize,
5342    xs: [&[i8]; 4],
5343    scales: &[f32],
5344) -> [f32; 4] {
5345    // SAFETY: as dot_q4tp_row_1x4_avx2; caller guarantees the ISA.
5346    unsafe {
5347        use core::arch::x86_64::*;
5348        let lomask = _mm256_set1_epi8(0x0F);
5349        let eight = _mm256_set1_epi8(8);
5350        let zero = _mm512_setzero_si512();
5351        let (mut v0, mut v1, mut v2, mut v3) = (
5352            _mm512_setzero_ps(),
5353            _mm512_setzero_ps(),
5354            _mm512_setzero_ps(),
5355            _mm512_setzero_ps(),
5356        );
5357        let pairs = gpr / 2;
5358        for gp in 0..pairs {
5359            let gi = gp * 2;
5360            let t = nib.as_ptr().add((r * gpr + gi) * Q4TP_NIB);
5361            let bb = _mm256_loadu_si256(t as *const __m256i);
5362            let lo = _mm256_and_si256(bb, lomask);
5363            let hi = _mm256_and_si256(_mm256_srli_epi16::<4>(bb), lomask);
5364            // `unpack` works per 128-bit lane, so the halves come out as
5365            // [A.lo, B.lo] and [A.hi, B.hi]; the shuffle reorders the four
5366            // 128-bit lanes into the weights' natural order, which is what
5367            // the straight activation load expects.
5368            let ul = _mm256_sub_epi8(_mm256_unpacklo_epi8(lo, hi), eight);
5369            let uh = _mm256_sub_epi8(_mm256_unpackhi_epi8(lo, hi), eight);
5370            let cat = _mm512_inserti64x4::<1>(_mm512_castsi256_si512(ul), uh);
5371            let w = _mm512_shuffle_i64x2::<0b11_01_10_00>(cat, cat);
5372            let wabs = _mm512_abs_epi8(w);
5373            let neg = _mm512_movepi8_mask(w);
5374            let off = gi * GROUP_SIZE;
5375            let sv = _mm512_insertf32x8::<1>(
5376                _mm512_castps256_ps512(_mm256_set1_ps(*scales.get_unchecked(gi))),
5377                _mm256_set1_ps(*scales.get_unchecked(gi + 1)),
5378            );
5379            let dot = |x: &[i8]| -> __m512 {
5380                let xv = _mm512_loadu_si512(x.as_ptr().add(off) as *const __m512i);
5381                let sx = _mm512_mask_sub_epi8(xv, neg, zero, xv);
5382                _mm512_cvtepi32_ps(_mm512_dpbusd_epi32(zero, wabs, sx))
5383            };
5384            v0 = _mm512_fmadd_ps(dot(xs[0]), sv, v0);
5385            v1 = _mm512_fmadd_ps(dot(xs[1]), sv, v1);
5386            v2 = _mm512_fmadd_ps(dot(xs[2]), sv, v2);
5387            v3 = _mm512_fmadd_ps(dot(xs[3]), sv, v3);
5388        }
5389        let mut acc = [
5390            _mm512_reduce_add_ps(v0),
5391            _mm512_reduce_add_ps(v1),
5392            _mm512_reduce_add_ps(v2),
5393            _mm512_reduce_add_ps(v3),
5394        ];
5395        // An odd group count leaves one group over; the narrow kernel
5396        // finishes it rather than the tail being a special case here.
5397        if gpr % 2 == 1 {
5398            let off = (gpr - 1) * GROUP_SIZE;
5399            for j in off..off + GROUP_SIZE {
5400                let (w, s) = q4tp_outlier(nib, r, gpr, j, scales);
5401                let ws = w * s;
5402                for k in 0..4 {
5403                    acc[k] += ws * *xs[k].get_unchecked(j) as f32;
5404                }
5405            }
5406        }
5407        acc
5408    }
5409}
5410
5411/// Four batch columns against one q4tp row: the tile is unpacked ONCE and
5412/// spent on four activation streams, which is where a prefill batch stops
5413/// being weight-bandwidth-bound. Twin of `dot_q4t_row_1x4_sdot`.
5414#[cfg(target_arch = "aarch64")]
5415#[target_feature(enable = "neon,dotprod")]
5416unsafe fn dot_q4tp_row_1x4_sdot(
5417    nib: &[u8],
5418    r: usize,
5419    gpr: usize,
5420    xs: [&[i8]; 4],
5421    scales: &[f32],
5422) -> [f32; 4] {
5423    // SAFETY: see dot_q4tp_row_sdot; every xs[k] is gpr·GROUP_SIZE long.
5424    unsafe {
5425        use core::arch::aarch64::*;
5426        use core::arch::asm;
5427        let lomask = vdupq_n_u8(0x0F);
5428        let eight = vdupq_n_s8(8);
5429        // Named accumulators, NOT an array indexed by a loop variable: the
5430        // latter does not stay in registers (the same defect cost 2x in the
5431        // AVX2 q4t kernel and again in WGSL).
5432        //
5433        // They are VECTORS, and the horizontal add happens once at the end
5434        // instead of once per group per column. `vaddvq` is a cross-lane
5435        // reduction — with 72 groups and four columns the old shape paid
5436        // 288 of them per row, each one a dependency stall the pipeline
5437        // cannot hide, to save four float adds. The group's scale now
5438        // rides an fma into the lane accumulators, so the arithmetic per
5439        // group is one convert and one fma. Summation order changes (the
5440        // lanes carry independent partial sums), which is the same
5441        // round-off class the SDOT path already lives in — the strict
5442        // kernel (`CMF_SDOT=0`, what `cortiq ppl` runs) is unchanged and
5443        // stays the reference.
5444        let (mut v0, mut v1, mut v2, mut v3) = (
5445            vdupq_n_f32(0.0),
5446            vdupq_n_f32(0.0),
5447            vdupq_n_f32(0.0),
5448            vdupq_n_f32(0.0),
5449        );
5450        for gi in 0..gpr {
5451            let t = nib.as_ptr().add((r * gpr + gi) * Q4TP_NIB);
5452            let s = *scales.get_unchecked(gi);
5453            let bb = vld1q_u8(t);
5454            let lo = vandq_u8(bb, lomask);
5455            let hi = vshrq_n_u8::<4>(bb);
5456            let e0 = vsubq_s8(vreinterpretq_s8_u8(vzip1q_u8(lo, hi)), eight);
5457            let e1 = vsubq_s8(vreinterpretq_s8_u8(vzip2q_u8(lo, hi)), eight);
5458            let off = gi * GROUP_SIZE;
5459            let dot4 = |x: &[i8]| -> int32x4_t {
5460                let x0 = vld1q_s8(x.as_ptr().add(off));
5461                let x1 = vld1q_s8(x.as_ptr().add(off + 16));
5462                let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
5463                asm!(
5464                    "sdot {a0:v}.4s, {e0:v}.16b, {x0:v}.16b",
5465                    "sdot {a1:v}.4s, {e1:v}.16b, {x1:v}.16b",
5466                    a0 = inout(vreg) a0, a1 = inout(vreg) a1,
5467                    e0 = in(vreg) e0, x0 = in(vreg) x0, e1 = in(vreg) e1, x1 = in(vreg) x1,
5468                    options(pure, nomem, nostack),
5469                );
5470                vaddq_s32(a0, a1)
5471            };
5472            v0 = vfmaq_n_f32(v0, vcvtq_f32_s32(dot4(xs[0])), s);
5473            v1 = vfmaq_n_f32(v1, vcvtq_f32_s32(dot4(xs[1])), s);
5474            v2 = vfmaq_n_f32(v2, vcvtq_f32_s32(dot4(xs[2])), s);
5475            v3 = vfmaq_n_f32(v3, vcvtq_f32_s32(dot4(xs[3])), s);
5476        }
5477        [
5478            vaddvq_f32(v0),
5479            vaddvq_f32(v1),
5480            vaddvq_f32(v2),
5481            vaddvq_f32(v3),
5482        ]
5483    }
5484}
5485
5486/// Fused q4tp matmat — the same three arms `q4t_matmat` has. Shipping only
5487/// the scalar one made Nanbeige-3B decode at 1.2 tok/s against q4t's 5.9:
5488/// the format was fine, the missing arms were the whole regression.
5489fn q4tp_matmat(
5490    bytes: &[u8],
5491    xs_all: &[f32],
5492    b: usize,
5493    rows: usize,
5494    cols: usize,
5495    out: &mut [f32],
5496    pool: Option<&Pool>,
5497) {
5498    debug_assert_eq!(out.len(), b * rows);
5499    let gpr = cols / GROUP_SIZE;
5500    let v = Q4tpView::new(bytes, rows, cols);
5501
5502    // Wide batches ride the AMX through a dequant-tile sgemm, as in q4t.
5503    #[cfg(target_os = "macos")]
5504    if b >= 8 && rows * cols >= 500_000 && accel_gemm_enabled() {
5505        dequant_matmat_accel(
5506            &|r, dst| {
5507                let mut sc = [0f32; 32];
5508                let mut scv;
5509                let s: &[f32] = if gpr <= 32 {
5510                    v.scales_into(r, gpr, &mut sc);
5511                    &sc[..gpr]
5512                } else {
5513                    scv = vec![0f32; gpr];
5514                    v.scales_into(r, gpr, &mut scv);
5515                    &scv
5516                };
5517                for gi in 0..gpr {
5518                    let tile = &v.nib[(r * gpr + gi) * Q4TP_NIB..(r * gpr + gi + 1) * Q4TP_NIB];
5519                    for (k, &bb) in tile.iter().enumerate() {
5520                        dst[gi * GROUP_SIZE + k * 2] = ((bb & 0x0F) as f32 - 8.0) * s[gi];
5521                        dst[gi * GROUP_SIZE + k * 2 + 1] =
5522                            (((bb >> 4) & 0x0F) as f32 - 8.0) * s[gi];
5523                    }
5524                }
5525            },
5526            xs_all,
5527            b,
5528            rows,
5529            cols,
5530            out,
5531            pool,
5532        );
5533        return;
5534    }
5535
5536    let out_addr = SendMut(out.as_mut_ptr());
5537    if a8w8_enabled() {
5538        let acts: Vec<SplitAct> = (0..b)
5539            .map(|bi| split_act(&xs_all[bi * cols..(bi + 1) * cols]))
5540            .collect();
5541        let acts = &acts;
5542        #[cfg(target_arch = "aarch64")]
5543        let blocked_ok = sdot_enabled() && blocked_enabled();
5544        // x86 gets the same blocking: one tile unpack spent on four
5545        // columns. Without it every column re-decoded the row, which is
5546        // why a 48-core EPYC measured a sixth of an M4's per-core rate.
5547        // The gate is `avx2_enabled`, as in q4t — `sdot_enabled` answers
5548        // for ARM's dotprod and is hard-wired false everywhere else, so
5549        // asking it here left the whole blocked path unreachable on x86.
5550        #[cfg(target_arch = "x86_64")]
5551        let blocked_ok = q4tp_blocked_x86();
5552        #[cfg(not(any(target_arch = "aarch64", target_arch = "x86_64")))]
5553        let blocked_ok = false;
5554        // Columns are swept in panels that fit L2. Without this a
5555        // row-pair walks every activation in the batch — 4.8 MB at
5556        // 512x512 — and does it again for the next pair, so the whole
5557        // batch streams out of the shared cache once per row. Measured
5558        // 800 GB/s of it, flat across batch sizes, which is the signature
5559        // of a loop bound by traffic rather than by arithmetic. A panel of
5560        // 256 columns is 590 KB beside 221 KB of this worker's weights:
5561        // both stay resident and the batch crosses L3 once instead of
5562        // once per row.
5563        let panel_cols: usize = std::env::var("CMF_Q4TP_PANEL")
5564            .ok()
5565            .and_then(|v| v.parse().ok())
5566            .filter(|v| *v > 0)
5567            .unwrap_or(256);
5568        let run = |start: usize, end: usize| {
5569            for abase in (0..acts.len()).step_by(panel_cols) {
5570                let alen = (acts.len() - abase).min(panel_cols);
5571                let mut sc = vec![0f32; gpr];
5572                #[cfg(target_arch = "x86_64")]
5573                let mut r_lo = start;
5574                #[cfg(target_arch = "x86_64")]
5575                if blocked_ok && alen >= 8 {
5576                    let mut sc1 = vec![0f32; gpr];
5577                    while r_lo + 2 <= end {
5578                        v.scales_into(r_lo, gpr, &mut sc);
5579                        v.scales_into(r_lo + 1, gpr, &mut sc1);
5580                        let mut bi = 0usize;
5581                        while bi + 8 <= alen {
5582                            let xs = [
5583                                acts[abase + bi].xq.as_slice(),
5584                                acts[abase + bi + 1].xq.as_slice(),
5585                                acts[abase + bi + 2].xq.as_slice(),
5586                                acts[abase + bi + 3].xq.as_slice(),
5587                                acts[abase + bi + 4].xq.as_slice(),
5588                                acts[abase + bi + 5].xq.as_slice(),
5589                                acts[abase + bi + 6].xq.as_slice(),
5590                                acts[abase + bi + 7].xq.as_slice(),
5591                            ];
5592                            let d = unsafe { dot_q4tp_2x8_avx512(v.nib, r_lo, gpr, xs, &sc, &sc1) };
5593                            for (row, dr, scr) in [(r_lo, &d[0], &sc), (r_lo + 1, &d[1], &sc1)] {
5594                                for k in 0..8 {
5595                                    let act = &acts[abase + bi + k];
5596                                    let mut acc = dr[k] * act.sx;
5597                                    for &(j, xv) in &act.outliers {
5598                                        let (w, s) = q4tp_outlier(v.nib, row, gpr, j, scr);
5599                                        acc += w * s * xv;
5600                                    }
5601                                    // SAFETY: disjoint (bi, r) cells per worker.
5602                                    unsafe { *out_addr.at((abase + bi + k) * rows + row) = acc };
5603                                }
5604                            }
5605                            bi += 8;
5606                        }
5607                        // Columns past the last group of eight, both rows —
5608                        // the same single-row kernel the tail below uses.
5609                        for row in [r_lo, r_lo + 1] {
5610                            let scr: &[f32] = if row == r_lo { &sc } else { &sc1 };
5611                            for b2 in bi..alen {
5612                                let act = &acts[abase + b2];
5613                                let xs4 = [
5614                                    act.xq.as_slice(),
5615                                    act.xq.as_slice(),
5616                                    act.xq.as_slice(),
5617                                    act.xq.as_slice(),
5618                                ];
5619                                let d =
5620                                    unsafe { dot_q4tp_row_1x4_avx512(v.nib, row, gpr, xs4, scr) };
5621                                let mut acc = d[0] * act.sx;
5622                                for &(j, xv) in &act.outliers {
5623                                    let (w, s) = q4tp_outlier(v.nib, row, gpr, j, scr);
5624                                    acc += w * s * xv;
5625                                }
5626                                // SAFETY: disjoint (bi, r) cells per worker.
5627                                unsafe { *out_addr.at((abase + b2) * rows + row) = acc };
5628                            }
5629                        }
5630                        r_lo += 2;
5631                    }
5632                }
5633                #[cfg(target_arch = "x86_64")]
5634                let row_start = r_lo;
5635                #[cfg(not(target_arch = "x86_64"))]
5636                let row_start = start;
5637                for r in row_start..end {
5638                    v.scales_into(r, gpr, &mut sc);
5639                    let mut bi = 0usize;
5640                    #[cfg(target_arch = "x86_64")]
5641                    if blocked_ok {
5642                        while bi + 8 <= alen {
5643                            let xs = [
5644                                acts[abase + bi].xq.as_slice(),
5645                                acts[abase + bi + 1].xq.as_slice(),
5646                                acts[abase + bi + 2].xq.as_slice(),
5647                                acts[abase + bi + 3].xq.as_slice(),
5648                                acts[abase + bi + 4].xq.as_slice(),
5649                                acts[abase + bi + 5].xq.as_slice(),
5650                                acts[abase + bi + 6].xq.as_slice(),
5651                                acts[abase + bi + 7].xq.as_slice(),
5652                            ];
5653                            let d = unsafe { dot_q4tp_row_1x8_avx512(v.nib, r, gpr, xs, &sc) };
5654                            for k in 0..8 {
5655                                let act = &acts[abase + bi + k];
5656                                let mut acc = d[k] * act.sx;
5657                                for &(j, xv) in &act.outliers {
5658                                    let (w, s) = q4tp_outlier(v.nib, r, gpr, j, &sc);
5659                                    acc += w * s * xv;
5660                                }
5661                                // SAFETY: disjoint (bi, r) cells per worker.
5662                                unsafe { *out_addr.at((abase + bi + k) * rows + r) = acc };
5663                            }
5664                            bi += 8;
5665                        }
5666                        while bi + 4 <= alen {
5667                            let xs = [
5668                                acts[abase + bi].xq.as_slice(),
5669                                acts[abase + bi + 1].xq.as_slice(),
5670                                acts[abase + bi + 2].xq.as_slice(),
5671                                acts[abase + bi + 3].xq.as_slice(),
5672                            ];
5673                            let d = unsafe { dot_q4tp_row_1x4_avx512(v.nib, r, gpr, xs, &sc) };
5674                            for k in 0..4 {
5675                                let act = &acts[abase + bi + k];
5676                                let mut acc = d[k] * act.sx;
5677                                for &(j, xv) in &act.outliers {
5678                                    let (w, s) = q4tp_outlier(v.nib, r, gpr, j, &sc);
5679                                    acc += w * s * xv;
5680                                }
5681                                // SAFETY: disjoint (bi, r) cells per worker.
5682                                unsafe { *out_addr.at((abase + bi + k) * rows + r) = acc };
5683                            }
5684                            bi += 4;
5685                        }
5686                    }
5687                    #[cfg(target_arch = "aarch64")]
5688                    if blocked_ok {
5689                        while bi + 4 <= alen {
5690                            let xs = [
5691                                acts[abase + bi].xq.as_slice(),
5692                                acts[abase + bi + 1].xq.as_slice(),
5693                                acts[abase + bi + 2].xq.as_slice(),
5694                                acts[abase + bi + 3].xq.as_slice(),
5695                            ];
5696                            let d = unsafe {
5697                                if q4tp_v1() {
5698                                    dot_q4tp_row_1x4_sdot_v1(v.nib, r, gpr, xs, &sc)
5699                                } else {
5700                                    dot_q4tp_row_1x4_sdot(v.nib, r, gpr, xs, &sc)
5701                                }
5702                            };
5703                            for k in 0..4 {
5704                                let act = &acts[abase + bi + k];
5705                                let mut acc = d[k] * act.sx;
5706                                for &(j, xv) in &act.outliers {
5707                                    let (w, s) = q4tp_outlier(v.nib, r, gpr, j, &sc);
5708                                    acc += w * s * xv;
5709                                }
5710                                // SAFETY: disjoint (bi, r) cells per worker.
5711                                unsafe { *out_addr.at((abase + bi + k) * rows + r) = acc };
5712                            }
5713                            bi += 4;
5714                        }
5715                    }
5716                    let _ = blocked_ok;
5717                    while bi < alen {
5718                        let act = &acts[abase + bi];
5719                        let mut acc = dot_q4tp_row_i8(v.nib, r, gpr, &act.xq, &sc) * act.sx;
5720                        for &(j, xv) in &act.outliers {
5721                            let (w, s) = q4tp_outlier(v.nib, r, gpr, j, &sc);
5722                            acc += w * s * xv;
5723                        }
5724                        // SAFETY: disjoint (bi, r) cells per worker range.
5725                        unsafe { *out_addr.at((abase + bi) * rows + r) = acc };
5726                        bi += 1;
5727                    }
5728                }
5729            }
5730        };
5731        dispatch_rows(pool, rows, &run);
5732        return;
5733    }
5734
5735    let run = |start: usize, end: usize| {
5736        let mut sc = vec![0f32; gpr];
5737        for r in start..end {
5738            v.scales_into(r, gpr, &mut sc);
5739            for bi in 0..b {
5740                let x = &xs_all[bi * cols..(bi + 1) * cols];
5741                // SAFETY: disjoint (bi, r) cells per worker range.
5742                unsafe { *out_addr.at(bi * rows + r) = q4tp_row_exact(v.nib, r, gpr, x, &sc) };
5743            }
5744        }
5745    };
5746    dispatch_rows(pool, rows, &run);
5747}
5748
5749/// Fused q4_tiled matvec (dispatch mirrors `q4matvec`).
5750fn q4t_matvec(
5751    bytes: &[u8],
5752    x: &[f32],
5753    rows: usize,
5754    cols: usize,
5755    out: &mut [f32],
5756    pool: Option<&Pool>,
5757) {
5758    debug_assert_eq!(out.len(), rows);
5759    let gpr = cols / GROUP_SIZE;
5760    let out_addr = SendMut(out.as_mut_ptr());
5761    if a8w8_enabled() {
5762        let act = split_act(x);
5763        let run = move |start: usize, end: usize| {
5764            for r in start..end {
5765                let mut acc = dot_q4t_row_i8(bytes, r, gpr, &act.xq) * act.sx;
5766                for &(j, xv) in &act.outliers {
5767                    let (w, s) = q4t_outlier(bytes, r, gpr, j);
5768                    acc += w * s * xv;
5769                }
5770                // SAFETY: disjoint row ranges per worker.
5771                unsafe { *out_addr.at(r) = acc };
5772            }
5773        };
5774        dispatch_rows(pool, rows, &run);
5775        return;
5776    }
5777    let run = move |start: usize, end: usize| {
5778        for r in start..end {
5779            // SAFETY: disjoint row ranges per worker.
5780            unsafe { *out_addr.at(r) = q4t_row_exact(bytes, r, gpr, x) };
5781        }
5782    };
5783    dispatch_rows(pool, rows, &run);
5784}
5785
5786/// Fused two-input q4_tiled matvec (weights read once per pair).
5787#[allow(clippy::too_many_arguments)]
5788fn q4t_matvec2(
5789    bytes: &[u8],
5790    x1: &[f32],
5791    x2: &[f32],
5792    rows: usize,
5793    cols: usize,
5794    o1: &mut [f32],
5795    o2: &mut [f32],
5796    pool: Option<&Pool>,
5797) {
5798    let gpr = cols / GROUP_SIZE;
5799    let p1 = SendMut(o1.as_mut_ptr());
5800    let p2 = SendMut(o2.as_mut_ptr());
5801    if a8w8_enabled() {
5802        let a1 = split_act(x1);
5803        let a2 = split_act(x2);
5804        let run = move |start: usize, end: usize| {
5805            for r in start..end {
5806                let mut v1 = dot_q4t_row_i8(bytes, r, gpr, &a1.xq) * a1.sx;
5807                let mut v2 = dot_q4t_row_i8(bytes, r, gpr, &a2.xq) * a2.sx;
5808                for &(j, xv) in &a1.outliers {
5809                    let (w, s) = q4t_outlier(bytes, r, gpr, j);
5810                    v1 += w * s * xv;
5811                }
5812                for &(j, xv) in &a2.outliers {
5813                    let (w, s) = q4t_outlier(bytes, r, gpr, j);
5814                    v2 += w * s * xv;
5815                }
5816                // SAFETY: disjoint row ranges per worker.
5817                unsafe {
5818                    *p1.at(r) = v1;
5819                    *p2.at(r) = v2;
5820                }
5821            }
5822        };
5823        dispatch_rows(pool, rows, &run);
5824        return;
5825    }
5826    let run = move |start: usize, end: usize| {
5827        for r in start..end {
5828            // SAFETY: disjoint row ranges per worker.
5829            unsafe {
5830                *p1.at(r) = q4t_row_exact(bytes, r, gpr, x1);
5831                *p2.at(r) = q4t_row_exact(bytes, r, gpr, x2);
5832            }
5833        }
5834    };
5835    dispatch_rows(pool, rows, &run);
5836}
5837
5838/// Batched q4_tiled matmat: each row's tiles stream once per microbatch.
5839#[allow(clippy::too_many_arguments)]
5840/// Prefill GEMM through Accelerate for group-quantized codecs: a
5841/// caller-supplied row dequantizer fills f32 tiles (pool-parallel) and
5842/// each tile rides the AMX with one sgemm — the generic sibling of
5843/// `qmatmat_accel` (q8). Numerics are f32-GEMM (tolerance class);
5844/// decode (b=1) never takes this path.
5845#[cfg(target_os = "macos")]
5846fn dequant_matmat_accel(
5847    dequant_row: &(dyn Fn(usize, &mut [f32]) + Sync),
5848    xs_all: &[f32],
5849    b: usize,
5850    rows: usize,
5851    cols: usize,
5852    out: &mut [f32],
5853    pool: Option<&Pool>,
5854) {
5855    const TR: usize = 2048;
5856    thread_local! {
5857        static WTILE: std::cell::RefCell<Vec<f32>> = const { std::cell::RefCell::new(Vec::new()) };
5858    }
5859    WTILE.with(|wt| {
5860        let mut wtile = wt.borrow_mut();
5861        wtile.resize(TR * cols, 0.0);
5862        let mut r0 = 0usize;
5863        while r0 < rows {
5864            let tr = TR.min(rows - r0);
5865            let wt_addr = SendMut(wtile.as_mut_ptr());
5866            let run = |start: usize, end: usize| {
5867                for r in start..end {
5868                    // SAFETY: workers cover disjoint r ranges.
5869                    let dst = unsafe { std::slice::from_raw_parts_mut(wt_addr.at(r * cols), cols) };
5870                    dequant_row(r0 + r, dst);
5871                }
5872            };
5873            dispatch_rows(pool, tr, &run);
5874            unsafe {
5875                accel_blas::cblas_sgemm(
5876                    101, // RowMajor
5877                    111, // NoTrans A
5878                    112, // Trans B
5879                    b as i32,
5880                    tr as i32,
5881                    cols as i32,
5882                    1.0,
5883                    xs_all.as_ptr(),
5884                    cols as i32,
5885                    wtile.as_ptr(),
5886                    cols as i32,
5887                    0.0,
5888                    out.as_mut_ptr().add(r0),
5889                    rows as i32,
5890                );
5891            }
5892            r0 += tr;
5893        }
5894    });
5895}
5896
5897fn q4t_matmat(
5898    bytes: &[u8],
5899    xs_all: &[f32],
5900    b: usize,
5901    rows: usize,
5902    cols: usize,
5903    out: &mut [f32],
5904    pool: Option<&Pool>,
5905) {
5906    debug_assert_eq!(out.len(), b * rows);
5907    let gpr = cols / GROUP_SIZE;
5908    // Wide batches ride the AMX like q8's qmatmat: on Apple silicon
5909    // the dequant-tile sgemm is an order above the SDOT row loop for
5910    // prefill shapes (imagegen DiT forwards are exactly this).
5911    #[cfg(target_os = "macos")]
5912    if b >= 8 && rows * cols >= 500_000 && accel_gemm_enabled() {
5913        dequant_matmat_accel(
5914            &|r, dst| {
5915                for gi in 0..gpr {
5916                    let tile = &bytes[(r * gpr + gi) * Q4_TILE..(r * gpr + gi + 1) * Q4_TILE];
5917                    let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
5918                    for (k, &bb) in tile[2..].iter().enumerate() {
5919                        dst[gi * GROUP_SIZE + k * 2] = ((bb & 0x0F) as f32 - 8.0) * s;
5920                        dst[gi * GROUP_SIZE + k * 2 + 1] = (((bb >> 4) & 0x0F) as f32 - 8.0) * s;
5921                    }
5922                }
5923            },
5924            xs_all,
5925            b,
5926            rows,
5927            cols,
5928            out,
5929            pool,
5930        );
5931        return;
5932    }
5933    let out_addr = SendMut(out.as_mut_ptr());
5934    if a8w8_enabled() {
5935        let acts: Vec<SplitAct> = (0..b)
5936            .map(|bi| split_act(&xs_all[bi * cols..(bi + 1) * cols]))
5937            .collect();
5938        let acts = &acts;
5939        #[cfg(target_arch = "x86_64")]
5940        let blocked_ok = avx2_enabled() && blocked_enabled();
5941        #[cfg(target_arch = "aarch64")]
5942        let blocked_ok = sdot_enabled() && blocked_enabled();
5943        #[cfg(not(any(target_arch = "x86_64", target_arch = "aarch64")))]
5944        let blocked_ok = false;
5945        let run = move |start: usize, end: usize| {
5946            for r in start..end {
5947                let mut bi = 0usize;
5948                #[cfg(target_arch = "aarch64")]
5949                if blocked_ok {
5950                    while bi + 4 <= acts.len() {
5951                        let xs = [
5952                            acts[bi].xq.as_slice(),
5953                            acts[bi + 1].xq.as_slice(),
5954                            acts[bi + 2].xq.as_slice(),
5955                            acts[bi + 3].xq.as_slice(),
5956                        ];
5957                        let d = unsafe { dot_q4t_row_1x4_sdot(bytes, r, gpr, xs) };
5958                        for k in 0..4 {
5959                            let act = &acts[bi + k];
5960                            let mut acc = d[k] * act.sx;
5961                            for &(j, xv) in &act.outliers {
5962                                let (w, sc) = q4t_outlier(bytes, r, gpr, j);
5963                                acc += w * sc * xv;
5964                            }
5965                            // SAFETY: disjoint (bi, r) cells per worker.
5966                            unsafe { *out_addr.at((bi + k) * rows + r) = acc };
5967                        }
5968                        bi += 4;
5969                    }
5970                }
5971                #[cfg(target_arch = "x86_64")]
5972                if blocked_ok {
5973                    while bi + 4 <= acts.len() {
5974                        let xs = [
5975                            acts[bi].xq.as_slice(),
5976                            acts[bi + 1].xq.as_slice(),
5977                            acts[bi + 2].xq.as_slice(),
5978                            acts[bi + 3].xq.as_slice(),
5979                        ];
5980                        let d = unsafe {
5981                            if vnni_tiles_enabled() {
5982                                dot_q4t_row_1x4_vnni(bytes, r, gpr, xs)
5983                            } else {
5984                                dot_q4t_row_1x4_avx2(bytes, r, gpr, xs)
5985                            }
5986                        };
5987                        for k in 0..4 {
5988                            let act = &acts[bi + k];
5989                            let mut acc = d[k] * act.sx;
5990                            for &(j, xv) in &act.outliers {
5991                                let (w, sc) = q4t_outlier(bytes, r, gpr, j);
5992                                acc += w * sc * xv;
5993                            }
5994                            // SAFETY: disjoint (bi, r) cells per worker.
5995                            unsafe { *out_addr.at((bi + k) * rows + r) = acc };
5996                        }
5997                        bi += 4;
5998                    }
5999                }
6000                let _ = blocked_ok;
6001                while bi < acts.len() {
6002                    let act = &acts[bi];
6003                    let mut acc = dot_q4t_row_i8(bytes, r, gpr, &act.xq) * act.sx;
6004                    for &(j, xv) in &act.outliers {
6005                        let (w, s) = q4t_outlier(bytes, r, gpr, j);
6006                        acc += w * s * xv;
6007                    }
6008                    // SAFETY: disjoint (bi, r) cells per worker range.
6009                    unsafe { *out_addr.at(bi * rows + r) = acc };
6010                    bi += 1;
6011                }
6012            }
6013        };
6014        dispatch_rows(pool, rows, &run);
6015        return;
6016    }
6017    let run = move |start: usize, end: usize| {
6018        for r in start..end {
6019            for bi in 0..b {
6020                let x = &xs_all[bi * cols..(bi + 1) * cols];
6021                // SAFETY: disjoint (bi, r) cells per worker range.
6022                unsafe { *out_addr.at(bi * rows + r) = q4t_row_exact(bytes, r, gpr, x) };
6023            }
6024        }
6025    };
6026    dispatch_rows(pool, rows, &run);
6027}
6028
6029// ── q1 (dtype 12): binary weights, [f16 scale][4B sign bits] per
6030// 32-group tile. The kernel family mirrors q4_tiled: one sequential
6031// stream of 6-byte tiles, per-tile integer dot × scale, exact outlier
6032// correction (A8W8 contract), exact scalar path under CMF_SDOT=0. ──
6033
6034/// Per-32-group sums of the quantized activation — the ±1 identity's
6035/// shared half: `dot = −2·sdot(mask, x) − gsum[g]`, computed ONCE per
6036/// matvec and reused by every row.
6037fn q1_group_sums(xq: &[i8], gpr: usize) -> Vec<i32> {
6038    (0..gpr)
6039        .map(|gi| {
6040            xq[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE]
6041                .iter()
6042                .map(|&v| v as i32)
6043                .sum()
6044        })
6045        .collect()
6046}
6047
6048/// One q1 row via the A8W8 int8 path — mask-SDOT on ARM (no ±1
6049/// expansion at all), scalar bit loop elsewhere (AVX2 queued with the
6050/// x86 pass).
6051#[inline]
6052#[allow(unreachable_code)]
6053/// AVX2 q1 row via the same ±1 identity as the ARM sdot kernel: the
6054/// sign bits expand to a {0, −1} byte mask through shuffle+cmpeq, the
6055/// masked activation sums through maddubs(1, x&mask), and
6056/// `dot = −(2·masked_sum + Σx_group)` — bit-identical integer math.
6057#[cfg(target_arch = "x86_64")]
6058#[target_feature(enable = "avx2")]
6059unsafe fn dot_q1_row_avx2(bytes: &[u8], r: usize, gpr: usize, xq: &[i8], gsum: &[i32]) -> f32 {
6060    // SAFETY: callers uphold the 6B-tile and xq/gsum length contracts.
6061    unsafe {
6062        use core::arch::x86_64::*;
6063        // Byte j of the mask must replicate bits-byte j/8.
6064        let expand = _mm256_setr_epi8(
6065            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,
6066            3, 3, 3,
6067        );
6068        let bitsel = _mm256_setr_epi8(
6069            1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64,
6070            -128, 1, 2, 4, 8, 16, 32, 64, -128,
6071        );
6072        let ones8 = _mm256_set1_epi8(1);
6073        let ones16 = _mm256_set1_epi16(1);
6074        let mut acc = 0f32;
6075        for gi in 0..gpr {
6076            let t = bytes.as_ptr().add((r * gpr + gi) * Q1_TILE);
6077            let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
6078            let bits = u32::from_le_bytes([*t.add(2), *t.add(3), *t.add(4), *t.add(5)]);
6079            let bc = _mm256_shuffle_epi8(_mm256_set1_epi32(bits as i32), expand);
6080            let mask = _mm256_cmpeq_epi8(_mm256_and_si256(bc, bitsel), bitsel);
6081            let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
6082            let sel = _mm256_and_si256(x, mask);
6083            // Σ of selected i8 lanes: maddubs(1u8, sel_i8) pairs → madd.
6084            let p16 = _mm256_maddubs_epi16(ones8, sel);
6085            let d32 = _mm256_madd_epi16(p16, ones16);
6086            let hi128 = _mm256_extracti128_si256::<1>(d32);
6087            let s128 = _mm_add_epi32(_mm256_castsi256_si128(d32), hi128);
6088            let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
6089            let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
6090            let msum = _mm_cvtsi128_si32(s32);
6091            // The and-select keeps x UN-negated (unlike ARM's −1-mask
6092            // sdot): d = Σ_set − Σ_unset = 2·Σ_set − Σ_all.
6093            let d = 2 * msum - gsum[gi];
6094            acc += d as f32 * s;
6095        }
6096        acc
6097    }
6098}
6099
6100/// VNNI twin of `dot_q1_row_avx2`: the masked-select sum goes through
6101/// one `vpdpbusd(1u8, sel)` (see `dpbusd_hsum` — bit-identical).
6102#[cfg(target_arch = "x86_64")]
6103#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
6104unsafe fn dot_q1_row_vnni(bytes: &[u8], r: usize, gpr: usize, xq: &[i8], gsum: &[i32]) -> f32 {
6105    // SAFETY: callers uphold the 6B-tile and xq/gsum length contracts.
6106    unsafe {
6107        use core::arch::x86_64::*;
6108        let expand = _mm256_setr_epi8(
6109            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,
6110            3, 3, 3,
6111        );
6112        let bitsel = _mm256_setr_epi8(
6113            1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64,
6114            -128, 1, 2, 4, 8, 16, 32, 64, -128,
6115        );
6116        let ones8 = _mm256_set1_epi8(1);
6117        let mut acc = 0f32;
6118        for gi in 0..gpr {
6119            let t = bytes.as_ptr().add((r * gpr + gi) * Q1_TILE);
6120            let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
6121            let bits = u32::from_le_bytes([*t.add(2), *t.add(3), *t.add(4), *t.add(5)]);
6122            let bc = _mm256_shuffle_epi8(_mm256_set1_epi32(bits as i32), expand);
6123            let mask = _mm256_cmpeq_epi8(_mm256_and_si256(bc, bitsel), bitsel);
6124            let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
6125            let msum = dpbusd_hsum(ones8, _mm256_and_si256(x, mask));
6126            let d = 2 * msum - gsum[gi];
6127            acc += d as f32 * s;
6128        }
6129        acc
6130    }
6131}
6132
6133/// VNNI twin of `dot_q1_row_1x4_avx2` (see `dpbusd_hsum`).
6134#[cfg(target_arch = "x86_64")]
6135#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
6136unsafe fn dot_q1_row_1x4_vnni(
6137    bytes: &[u8],
6138    r: usize,
6139    gpr: usize,
6140    xs: [&[i8]; 4],
6141    gsums: [&[i32]; 4],
6142) -> [f32; 4] {
6143    // SAFETY: callers uphold the 6B-tile and xq/gsum length contracts.
6144    unsafe {
6145        use core::arch::x86_64::*;
6146        let expand = _mm256_setr_epi8(
6147            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,
6148            3, 3, 3,
6149        );
6150        let bitsel = _mm256_setr_epi8(
6151            1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64,
6152            -128, 1, 2, 4, 8, 16, 32, 64, -128,
6153        );
6154        let ones8 = _mm256_set1_epi8(1);
6155        let mut acc = [0f32; 4];
6156        for gi in 0..gpr {
6157            let t = bytes.as_ptr().add((r * gpr + gi) * Q1_TILE);
6158            let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
6159            let bits = u32::from_le_bytes([*t.add(2), *t.add(3), *t.add(4), *t.add(5)]);
6160            let bc = _mm256_shuffle_epi8(_mm256_set1_epi32(bits as i32), expand);
6161            let mask = _mm256_cmpeq_epi8(_mm256_and_si256(bc, bitsel), bitsel);
6162            for (k, xq) in xs.iter().enumerate() {
6163                let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
6164                let msum = dpbusd_hsum(ones8, _mm256_and_si256(x, mask));
6165                let d = 2 * msum - gsums[k][gi];
6166                acc[k] += d as f32 * s;
6167            }
6168        }
6169        acc
6170    }
6171}
6172
6173/// The blocked 1×4 flavor: the expanded bit mask serves four activation
6174/// streams per group (mask build once, four select+reduce chains).
6175#[cfg(target_arch = "x86_64")]
6176#[target_feature(enable = "avx2")]
6177unsafe fn dot_q1_row_1x4_avx2(
6178    bytes: &[u8],
6179    r: usize,
6180    gpr: usize,
6181    xs: [&[i8]; 4],
6182    gsums: [&[i32]; 4],
6183) -> [f32; 4] {
6184    // SAFETY: callers uphold the 6B-tile and xq/gsum length contracts.
6185    unsafe {
6186        use core::arch::x86_64::*;
6187        let expand = _mm256_setr_epi8(
6188            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,
6189            3, 3, 3,
6190        );
6191        let bitsel = _mm256_setr_epi8(
6192            1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64,
6193            -128, 1, 2, 4, 8, 16, 32, 64, -128,
6194        );
6195        let ones8 = _mm256_set1_epi8(1);
6196        let ones16 = _mm256_set1_epi16(1);
6197        let mut acc = [0f32; 4];
6198        for gi in 0..gpr {
6199            let t = bytes.as_ptr().add((r * gpr + gi) * Q1_TILE);
6200            let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
6201            let bits = u32::from_le_bytes([*t.add(2), *t.add(3), *t.add(4), *t.add(5)]);
6202            let bc = _mm256_shuffle_epi8(_mm256_set1_epi32(bits as i32), expand);
6203            let mask = _mm256_cmpeq_epi8(_mm256_and_si256(bc, bitsel), bitsel);
6204            for (k, xq) in xs.iter().enumerate() {
6205                let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
6206                let sel = _mm256_and_si256(x, mask);
6207                let p16 = _mm256_maddubs_epi16(ones8, sel);
6208                let d32 = _mm256_madd_epi16(p16, ones16);
6209                let hi128 = _mm256_extracti128_si256::<1>(d32);
6210                let s128 = _mm_add_epi32(_mm256_castsi256_si128(d32), hi128);
6211                let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
6212                let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
6213                let msum = _mm_cvtsi128_si32(s32);
6214                let d = 2 * msum - gsums[k][gi];
6215                acc[k] += d as f32 * s;
6216            }
6217        }
6218        acc
6219    }
6220}
6221
6222#[allow(unreachable_code)]
6223fn dot_q1_row_i8(bytes: &[u8], r: usize, gpr: usize, xq: &[i8], gsum: &[i32]) -> f32 {
6224    #[cfg(target_arch = "aarch64")]
6225    unsafe {
6226        return dot_q1_row_sdot(bytes, r, gpr, xq, gsum);
6227    }
6228    #[cfg(target_arch = "x86_64")]
6229    if avx2_enabled() {
6230        unsafe {
6231            if vnni_tiles_enabled() {
6232                return dot_q1_row_vnni(bytes, r, gpr, xq, gsum);
6233            }
6234            return dot_q1_row_avx2(bytes, r, gpr, xq, gsum);
6235        }
6236    }
6237    let _ = gsum;
6238    let mut acc = 0f32;
6239    for gi in 0..gpr {
6240        let tile = &bytes[(r * gpr + gi) * Q1_TILE..(r * gpr + gi + 1) * Q1_TILE];
6241        let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
6242        let mut d = 0i32;
6243        for (j, &b) in tile[2..].iter().enumerate() {
6244            for k in 0..8 {
6245                let w = ((b >> k) & 1) as i32 * 2 - 1;
6246                d += w * xq[gi * GROUP_SIZE + j * 8 + k] as i32;
6247            }
6248        }
6249        acc += d as f32 * s;
6250    }
6251    acc
6252}
6253
6254/// SDOT q1 row via the ±1 identity: the vtst mask (0xFF where the bit
6255/// is set, i.e. −1 as i8) feeds `sdot` DIRECTLY — no expansion to ±1
6256/// lanes at all — and `dot = −(2·sdot(mask, x) + Σx_group)`, with the
6257/// per-group activation sums shared across every row of the matvec.
6258/// Four tiles (128 weights) per iteration: integer dots reduce through
6259/// a vpaddq tree into ONE i32x4 that meets its four scales in a single
6260/// fused f32 multiply-add. Integer math throughout — bit-identical to
6261/// the scalar ±1 reference.
6262#[cfg(target_arch = "aarch64")]
6263#[target_feature(enable = "neon,dotprod")]
6264unsafe fn dot_q1_row_sdot(bytes: &[u8], r: usize, gpr: usize, xq: &[i8], gsum: &[i32]) -> f32 {
6265    // SAFETY: callers uphold slice-length contracts (6B tile per group,
6266    // xq.len() == gpr·GROUP_SIZE, gsum.len() == gpr).
6267    unsafe {
6268        use core::arch::aarch64::*;
6269        use core::arch::asm;
6270        const MASKS: [u8; 16] = [1, 2, 4, 8, 16, 32, 64, 128, 1, 2, 4, 8, 16, 32, 64, 128];
6271        let m = vld1q_u8(MASKS.as_ptr());
6272        // One tile's −Σ_set(x) as an UNREDUCED i32x4 (two mask-sdots).
6273        macro_rules! tile_dot {
6274            ($t:expr, $x:expr) => {{
6275                let v0 = vcombine_u8(vdup_n_u8(*$t.add(2)), vdup_n_u8(*$t.add(3)));
6276                let v1 = vcombine_u8(vdup_n_u8(*$t.add(4)), vdup_n_u8(*$t.add(5)));
6277                let w0 = vreinterpretq_s8_u8(vtstq_u8(v0, m));
6278                let w1 = vreinterpretq_s8_u8(vtstq_u8(v1, m));
6279                let x0 = vld1q_s8($x);
6280                let x1 = vld1q_s8($x.add(16));
6281                let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
6282                asm!(
6283                    "sdot {a0:v}.4s, {w0:v}.16b, {x0:v}.16b",
6284                    "sdot {a1:v}.4s, {w1:v}.16b, {x1:v}.16b",
6285                    a0 = inout(vreg) a0, a1 = inout(vreg) a1,
6286                    w0 = in(vreg) w0, x0 = in(vreg) x0, w1 = in(vreg) w1, x1 = in(vreg) x1,
6287                    options(pure, nomem, nostack),
6288                );
6289                vaddq_s32(a0, a1)
6290            }};
6291        }
6292        // TBL unpack over PAIR loads: one vld1q covers two 6B tiles
6293        // ([s s b b b b][s s b b b b] + 4B slack), TBL replicates each
6294        // bit-byte across 8 lanes for vtst, and the four scales gather
6295        // through tbl2 into one fcvtl — the 16 ld1r broadcast loads and
6296        // 4 branchy software f16 conversions per 128 weights (the
6297        // measured load-port wall of this kernel) become 2 vector
6298        // loads + 9 table lookups. Integer math order is unchanged —
6299        // bit-identical results (FCVTL is exact on every f16).
6300        const IW00: [u8; 16] = [2, 2, 2, 2, 2, 2, 2, 2, 3, 3, 3, 3, 3, 3, 3, 3];
6301        const IW01: [u8; 16] = [4, 4, 4, 4, 4, 4, 4, 4, 5, 5, 5, 5, 5, 5, 5, 5];
6302        const IW10: [u8; 16] = [8, 8, 8, 8, 8, 8, 8, 8, 9, 9, 9, 9, 9, 9, 9, 9];
6303        const IW11: [u8; 16] = [
6304            10, 10, 10, 10, 10, 10, 10, 10, 11, 11, 11, 11, 11, 11, 11, 11,
6305        ];
6306        const ISC: [u8; 8] = [0, 1, 6, 7, 16, 17, 22, 23];
6307        let (iw00, iw01) = (vld1q_u8(IW00.as_ptr()), vld1q_u8(IW01.as_ptr()));
6308        let (iw10, iw11) = (vld1q_u8(IW10.as_ptr()), vld1q_u8(IW11.as_ptr()));
6309        let isc = vld1_u8(ISC.as_ptr());
6310        // One tile's −Σ_set(x) from a TBL-unpacked pair load.
6311        macro_rules! tile_dot_tbl {
6312            ($ld:expr, $i0:expr, $i1:expr, $x:expr) => {{
6313                let w0 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8($ld, $i0), m));
6314                let w1 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8($ld, $i1), m));
6315                let x0 = vld1q_s8($x);
6316                let x1 = vld1q_s8($x.add(16));
6317                let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
6318                asm!(
6319                    "sdot {a0:v}.4s, {w0:v}.16b, {x0:v}.16b",
6320                    "sdot {a1:v}.4s, {w1:v}.16b, {x1:v}.16b",
6321                    a0 = inout(vreg) a0, a1 = inout(vreg) a1,
6322                    w0 = in(vreg) w0, x0 = in(vreg) x0, w1 = in(vreg) w1, x1 = in(vreg) x1,
6323                    options(pure, nomem, nostack),
6324                );
6325                vaddq_s32(a0, a1)
6326            }};
6327        }
6328        let base = bytes.as_ptr().add(r * gpr * Q1_TILE);
6329        let row_base = r * gpr * Q1_TILE;
6330        let abs_end = bytes.len();
6331        let xp = xq.as_ptr();
6332        let gp = gsum.as_ptr();
6333        let mut accv = vdupq_n_f32(0.0);
6334        let mut gi = 0;
6335        // The second pair load reads 4B past tile gi+3 — stay inside
6336        // the payload slice (only the file's final tiles fall back).
6337        while gi + 4 <= gpr && row_base + (gi + 4) * Q1_TILE + 4 <= abs_end {
6338            let t0 = base.add(gi * Q1_TILE);
6339            let ld_a = vld1q_u8(t0);
6340            let ld_b = vld1q_u8(t0.add(2 * Q1_TILE));
6341            let d0 = tile_dot_tbl!(ld_a, iw00, iw01, xp.add(gi * GROUP_SIZE));
6342            let d1 = tile_dot_tbl!(ld_a, iw10, iw11, xp.add((gi + 1) * GROUP_SIZE));
6343            let d2 = tile_dot_tbl!(ld_b, iw00, iw01, xp.add((gi + 2) * GROUP_SIZE));
6344            let d3 = tile_dot_tbl!(ld_b, iw10, iw11, xp.add((gi + 3) * GROUP_SIZE));
6345            // [−Σ0, −Σ1, −Σ2, −Σ3] → dots = −(2·Σset_neg + gsum)
6346            let neg = vpaddq_s32(vpaddq_s32(d0, d1), vpaddq_s32(d2, d3));
6347            let g = vld1q_s32(gp.add(gi));
6348            let dots = vnegq_s32(vaddq_s32(vshlq_n_s32::<1>(neg), g));
6349            let sc16 = vqtbl2_u8(uint8x16x2_t(ld_a, ld_b), isc);
6350            let scf: float32x4_t;
6351            asm!(
6352                "fcvtl {o:v}.4s, {i:v}.4h",
6353                o = out(vreg) scf, i = in(vreg) sc16,
6354                options(pure, nomem, nostack),
6355            );
6356            accv = vfmaq_f32(accv, vcvtq_f32_s32(dots), scf);
6357            gi += 4;
6358        }
6359        let mut acc = vaddvq_f32(accv);
6360        while gi < gpr {
6361            let t = base.add(gi * Q1_TILE);
6362            let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
6363            let d = vaddvq_s32(tile_dot!(t, xp.add(gi * GROUP_SIZE)));
6364            acc += (-(2 * d + *gp.add(gi))) as f32 * s;
6365            gi += 1;
6366        }
6367        acc
6368    }
6369}
6370
6371/// Blocked q1 1×4: one TBL unpack of the tile pair serves FOUR
6372/// activation streams (prefill amortization — the same idea as the
6373/// AVX2 twin; per stream the group order, fma order and tail match the
6374/// single-row kernel exactly, so batch == matvec bit-for-bit).
6375#[cfg(target_arch = "aarch64")]
6376#[target_feature(enable = "neon,dotprod")]
6377unsafe fn dot_q1_row_1x4_sdot(
6378    bytes: &[u8],
6379    r: usize,
6380    gpr: usize,
6381    xs: [&[i8]; 4],
6382    gs: [&[i32]; 4],
6383) -> [f32; 4] {
6384    // SAFETY: same slice-length contracts as `dot_q1_row_sdot`, ×4.
6385    unsafe {
6386        use core::arch::aarch64::*;
6387        use core::arch::asm;
6388        const MASKS: [u8; 16] = [1, 2, 4, 8, 16, 32, 64, 128, 1, 2, 4, 8, 16, 32, 64, 128];
6389        const IW00: [u8; 16] = [2, 2, 2, 2, 2, 2, 2, 2, 3, 3, 3, 3, 3, 3, 3, 3];
6390        const IW01: [u8; 16] = [4, 4, 4, 4, 4, 4, 4, 4, 5, 5, 5, 5, 5, 5, 5, 5];
6391        const IW10: [u8; 16] = [8, 8, 8, 8, 8, 8, 8, 8, 9, 9, 9, 9, 9, 9, 9, 9];
6392        const IW11: [u8; 16] = [
6393            10, 10, 10, 10, 10, 10, 10, 10, 11, 11, 11, 11, 11, 11, 11, 11,
6394        ];
6395        const ISC: [u8; 8] = [0, 1, 6, 7, 16, 17, 22, 23];
6396        let m = vld1q_u8(MASKS.as_ptr());
6397        let (iw00, iw01) = (vld1q_u8(IW00.as_ptr()), vld1q_u8(IW01.as_ptr()));
6398        let (iw10, iw11) = (vld1q_u8(IW10.as_ptr()), vld1q_u8(IW11.as_ptr()));
6399        let isc = vld1_u8(ISC.as_ptr());
6400        macro_rules! sdot2 {
6401            ($w0:expr, $w1:expr, $x:expr) => {{
6402                let x0 = vld1q_s8($x);
6403                let x1 = vld1q_s8($x.add(16));
6404                let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
6405                asm!(
6406                    "sdot {a0:v}.4s, {w0:v}.16b, {x0:v}.16b",
6407                    "sdot {a1:v}.4s, {w1:v}.16b, {x1:v}.16b",
6408                    a0 = inout(vreg) a0, a1 = inout(vreg) a1,
6409                    w0 = in(vreg) $w0, x0 = in(vreg) x0, w1 = in(vreg) $w1, x1 = in(vreg) x1,
6410                    options(pure, nomem, nostack),
6411                );
6412                vaddq_s32(a0, a1)
6413            }};
6414        }
6415        let base = bytes.as_ptr().add(r * gpr * Q1_TILE);
6416        let row_base = r * gpr * Q1_TILE;
6417        let abs_end = bytes.len();
6418        let mut accv = [vdupq_n_f32(0.0); 4];
6419        let mut gi = 0;
6420        while gi + 4 <= gpr && row_base + (gi + 4) * Q1_TILE + 4 <= abs_end {
6421            let t0 = base.add(gi * Q1_TILE);
6422            let ld_a = vld1q_u8(t0);
6423            let ld_b = vld1q_u8(t0.add(2 * Q1_TILE));
6424            // Unpack ONCE — eight ±mask vectors serve all four streams.
6425            let w00 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8(ld_a, iw00), m));
6426            let w01 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8(ld_a, iw01), m));
6427            let w10 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8(ld_a, iw10), m));
6428            let w11 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8(ld_a, iw11), m));
6429            let w20 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8(ld_b, iw00), m));
6430            let w21 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8(ld_b, iw01), m));
6431            let w30 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8(ld_b, iw10), m));
6432            let w31 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8(ld_b, iw11), m));
6433            let sc16 = vqtbl2_u8(uint8x16x2_t(ld_a, ld_b), isc);
6434            let scf: float32x4_t;
6435            asm!(
6436                "fcvtl {o:v}.4s, {i:v}.4h",
6437                o = out(vreg) scf, i = in(vreg) sc16,
6438                options(pure, nomem, nostack),
6439            );
6440            for k in 0..4 {
6441                let xp = xs[k].as_ptr();
6442                let d0 = sdot2!(w00, w01, xp.add(gi * GROUP_SIZE));
6443                let d1 = sdot2!(w10, w11, xp.add((gi + 1) * GROUP_SIZE));
6444                let d2 = sdot2!(w20, w21, xp.add((gi + 2) * GROUP_SIZE));
6445                let d3 = sdot2!(w30, w31, xp.add((gi + 3) * GROUP_SIZE));
6446                let neg = vpaddq_s32(vpaddq_s32(d0, d1), vpaddq_s32(d2, d3));
6447                let g = vld1q_s32(gs[k].as_ptr().add(gi));
6448                let dots = vnegq_s32(vaddq_s32(vshlq_n_s32::<1>(neg), g));
6449                accv[k] = vfmaq_f32(accv[k], vcvtq_f32_s32(dots), scf);
6450            }
6451            gi += 4;
6452        }
6453        let mut acc = [
6454            vaddvq_f32(accv[0]),
6455            vaddvq_f32(accv[1]),
6456            vaddvq_f32(accv[2]),
6457            vaddvq_f32(accv[3]),
6458        ];
6459        while gi < gpr {
6460            let t = base.add(gi * Q1_TILE);
6461            let sc = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
6462            let v0 = vcombine_u8(vdup_n_u8(*t.add(2)), vdup_n_u8(*t.add(3)));
6463            let v1 = vcombine_u8(vdup_n_u8(*t.add(4)), vdup_n_u8(*t.add(5)));
6464            let w0 = vreinterpretq_s8_u8(vtstq_u8(v0, m));
6465            let w1 = vreinterpretq_s8_u8(vtstq_u8(v1, m));
6466            for k in 0..4 {
6467                let d = vaddvq_s32(sdot2!(w0, w1, xs[k].as_ptr().add(gi * GROUP_SIZE)));
6468                acc[k] += (-(2 * d + *gs[k].as_ptr().add(gi))) as f32 * sc;
6469            }
6470            gi += 1;
6471        }
6472        acc
6473    }
6474}
6475
6476/// (weight ±1, scale) of one q1 element — the exact outlier term.
6477#[inline]
6478fn q1_outlier(bytes: &[u8], r: usize, gpr: usize, j: usize) -> (f32, f32) {
6479    let gi = j / GROUP_SIZE;
6480    let k = j % GROUP_SIZE;
6481    let tile = &bytes[(r * gpr + gi) * Q1_TILE..(r * gpr + gi + 1) * Q1_TILE];
6482    let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
6483    let bit = (tile[2 + k / 8] >> (k % 8)) & 1;
6484    ((bit as i32 * 2 - 1) as f32, s)
6485}
6486
6487/// Exact scalar q1 row (CMF_SDOT=0 contract).
6488#[inline]
6489fn q1_row_exact(bytes: &[u8], r: usize, gpr: usize, x: &[f32]) -> f32 {
6490    let mut acc = 0f32;
6491    for gi in 0..gpr {
6492        let tile = &bytes[(r * gpr + gi) * Q1_TILE..(r * gpr + gi + 1) * Q1_TILE];
6493        let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
6494        let xg = &x[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
6495        let mut ga = 0f32;
6496        for (j, &b) in tile[2..].iter().enumerate() {
6497            for k in 0..8 {
6498                ga += (((b >> k) & 1) as f32 * 2.0 - 1.0) * xg[j * 8 + k];
6499            }
6500        }
6501        acc += ga * s;
6502    }
6503    acc
6504}
6505
6506/// One q1 row range via A8W8 (the body of `q1_matvec`'s hot loop,
6507/// extracted so multi-matrix jobs drive the same kernel).
6508#[allow(clippy::too_many_arguments)]
6509fn q1_range_a8w8(
6510    bytes: &[u8],
6511    gpr: usize,
6512    act: &SplitAct,
6513    gsum: &[i32],
6514    out: SendMut,
6515    start: usize,
6516    end: usize,
6517) {
6518    for r in start..end {
6519        let mut acc = dot_q1_row_i8(bytes, r, gpr, &act.xq, gsum) * act.sx;
6520        for &(j, xv) in &act.outliers {
6521            let (w, s) = q1_outlier(bytes, r, gpr, j);
6522            acc += w * s * xv;
6523        }
6524        // SAFETY: disjoint row ranges per worker.
6525        unsafe { *out.at(r) = acc };
6526    }
6527}
6528
6529/// Exact-scalar q1 row range (CMF_SDOT=0 contract).
6530fn q1_range_f32(bytes: &[u8], gpr: usize, x: &[f32], out: SendMut, start: usize, end: usize) {
6531    for r in start..end {
6532        // SAFETY: disjoint row ranges per worker.
6533        unsafe { *out.at(r) = q1_row_exact(bytes, r, gpr, x) };
6534    }
6535}
6536
6537/// q1t per-row overlay locator. After the base (`base_len`) come
6538/// `[u32 row_ptr[rows+1]]` then `[(u16 col, f16 val)]` grouped by row (row
6539/// `r`'s entries are `[row_ptr[r], row_ptr[r+1])`). Returns
6540/// `(row_ptr offset, entries offset, present)`.
6541fn q1t_overlay(bytes: &[u8], base_len: usize, rows: usize) -> (usize, usize, bool) {
6542    let entries = base_len + (rows + 1) * 4;
6543    (base_len, entries, entries <= bytes.len())
6544}
6545
6546/// Read `row_ptr[r]` from the overlay's prefix-sum table.
6547#[inline]
6548fn q1t_rowptr(bytes: &[u8], rp_off: usize, r: usize) -> usize {
6549    let o = rp_off + r * 4;
6550    u32::from_le_bytes([bytes[o], bytes[o + 1], bytes[o + 2], bytes[o + 3]]) as usize
6551}
6552
6553/// Byte → the 5 ternary signs it packs `{−1,0,+1}` as f32, precomputed so
6554/// decoding a q1t code is a table load, not the base-3 divide/modulo per
6555/// weight (division is ~20–40× the cost of a load). Built at compile time.
6556const SIGN5: [[f32; 5]; 256] = {
6557    let mut lut = [[0.0f32; 5]; 256];
6558    let pow3 = [1u16, 3, 9, 27, 81];
6559    let mut byte = 0usize;
6560    while byte < 256 {
6561        let mut i = 0usize;
6562        while i < 5 {
6563            let code = (byte as u16 / pow3[i]) % 3;
6564            lut[byte][i] = if code == 1 {
6565                1.0
6566            } else if code == 2 {
6567                -1.0
6568            } else {
6569                0.0
6570            };
6571            i += 1;
6572        }
6573        byte += 1;
6574    }
6575    lut
6576};
6577
6578/// Same table, as i8 signs — the operand for the int8 SDOT base kernel.
6579const SIGN5_I8: [[i8; 5]; 256] = {
6580    let mut lut = [[0i8; 5]; 256];
6581    let pow3 = [1u16, 3, 9, 27, 81];
6582    let mut byte = 0usize;
6583    while byte < 256 {
6584        let mut i = 0usize;
6585        while i < 5 {
6586            let code = (byte as u16 / pow3[i]) % 3;
6587            lut[byte][i] = if code == 1 {
6588                1
6589            } else if code == 2 {
6590                -1
6591            } else {
6592                0
6593            };
6594            i += 1;
6595        }
6596        byte += 1;
6597    }
6598    lut
6599};
6600
6601/// The same 5 i8 signs packed into a u64 (`[s0 s1 s2 s3 s4 0 0 0]`, LE) so the
6602/// group unpack is 7 unaligned u64 stores at offsets 0,5,10,…,30 instead of
6603/// six 5-byte copies + LUT indexing — each store's trailing zeros are fixed by
6604/// the next store, and the last one runs 6 B past the 32nd weight (the unpack
6605/// buffer is padded to 40). This is the decode/prefill hot inner op.
6606const SIGN5_U64: [u64; 256] = {
6607    let mut lut = [0u64; 256];
6608    let pow3 = [1u16, 3, 9, 27, 81];
6609    let mut byte = 0usize;
6610    while byte < 256 {
6611        let mut v = 0u64;
6612        let mut i = 0usize;
6613        while i < 5 {
6614            let code = (byte as u16 / pow3[i]) % 3;
6615            let s: u8 = if code == 1 {
6616                1
6617            } else if code == 2 {
6618                0xFF
6619            } else {
6620                0
6621            };
6622            v |= (s as u64) << (i * 8);
6623            i += 1;
6624        }
6625        lut[byte] = v;
6626        byte += 1;
6627    }
6628    lut
6629};
6630
6631/// Ternary base weight at `(row r, col j)` = `sign(code)·s_group`. Used to add
6632/// back activation-outlier columns, whose `x` was zeroed for the int8 bulk dot
6633/// (`split_act`). At a weight-outlier position the code is 0, so this is 0 and
6634/// the overlay correction owns that column — no double counting.
6635#[inline]
6636fn q1t_base_weight(bytes: &[u8], r: usize, gpr: usize, j: usize) -> f32 {
6637    const TILE: usize = cortiq_core::quant::Q1T_TILE;
6638    let off = (r * gpr + j / GROUP_SIZE) * TILE;
6639    let s = f16_to_f32(u16::from_le_bytes([bytes[off], bytes[off + 1]]));
6640    let within = j % GROUP_SIZE;
6641    SIGN5[bytes[off + 2 + within / 5] as usize][within % 5] * s
6642}
6643
6644/// One 32-group int8 dot via two SDOTs. Bit-exact vs the scalar i8 sum
6645/// (integer accumulation is order-independent).
6646#[cfg(target_arch = "aarch64")]
6647#[target_feature(enable = "neon,dotprod")]
6648#[inline]
6649unsafe fn sdot32_i8(w: *const i8, x: *const i8) -> i32 {
6650    // SAFETY: caller guarantees 32 readable i8 at each pointer.
6651    unsafe {
6652        use core::arch::aarch64::*;
6653        use core::arch::asm;
6654        let w0 = vld1q_s8(w);
6655        let w1 = vld1q_s8(w.add(16));
6656        let x0 = vld1q_s8(x);
6657        let x1 = vld1q_s8(x.add(16));
6658        let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
6659        asm!(
6660            "sdot {a0:v}.4s, {w0:v}.16b, {x0:v}.16b",
6661            "sdot {a1:v}.4s, {w1:v}.16b, {x1:v}.16b",
6662            a0 = inout(vreg) a0, a1 = inout(vreg) a1,
6663            w0 = in(vreg) w0, x0 = in(vreg) x0, w1 = in(vreg) w1, x1 = in(vreg) x1,
6664            options(pure, nomem, nostack),
6665        );
6666        vaddvq_s32(vaddq_s32(a0, a1))
6667    }
6668}
6669
6670/// One 32-group int8 dot via AVX2: signed·signed as `maddubs(|w|, sign(x,w))`
6671/// then `madd` and a horizontal reduce (the same idiom as `dot_q4t_row_avx2`).
6672#[cfg(target_arch = "x86_64")]
6673#[target_feature(enable = "avx2")]
6674#[inline]
6675unsafe fn i8dot32_avx2(w: *const i8, x: *const i8) -> i32 {
6676    // SAFETY: caller guarantees 32 readable i8 at each pointer.
6677    unsafe {
6678        use core::arch::x86_64::*;
6679        let wv = _mm256_loadu_si256(w as *const __m256i);
6680        let xv = _mm256_loadu_si256(x as *const __m256i);
6681        let p16 = _mm256_maddubs_epi16(_mm256_abs_epi8(wv), _mm256_sign_epi8(xv, wv));
6682        let d = _mm256_madd_epi16(p16, _mm256_set1_epi16(1));
6683        let hi128 = _mm256_extracti128_si256::<1>(d);
6684        let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
6685        let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
6686        let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
6687        _mm_cvtsi128_si32(s32)
6688    }
6689}
6690
6691/// Unpack one q1t group's base-3 codes into 32 i8 signs via 7 unaligned u64
6692/// stores (see `SIGN5_U64`). `dst` MUST have ≥ 40 bytes: the 7th store writes
6693/// `dst[30..38]`. Stores go in order so each one's trailing zeros are
6694/// overwritten by the next; the final 6 padding bytes are unused by the dot.
6695#[inline]
6696fn q1t_unpack_group_i8(codes: *const u8, dst: &mut [i8]) {
6697    debug_assert!(dst.len() >= 40);
6698    // SAFETY: codes points at 7 readable bytes; dst has ≥ 40 bytes so every
6699    // 8-byte store at offset bi*5 (bi ≤ 6 → ≤ 30) stays in bounds.
6700    unsafe {
6701        let p = dst.as_mut_ptr();
6702        for bi in 0..7 {
6703            core::ptr::write_unaligned(
6704                p.add(bi * 5) as *mut u64,
6705                SIGN5_U64[*codes.add(bi) as usize],
6706            );
6707        }
6708    }
6709}
6710
6711/// One 32-group int8 dot, arch-dispatched (the matmat inner loop, where the
6712/// row's signs are unpacked once and dotted against every batch input).
6713/// Callers are gated by `a8w8_enabled()`, so the target-feature arms are
6714/// reachable; the scalar arm is a non-SIMD-arch fallback.
6715#[inline]
6716fn q1t_i8dot32(w: *const i8, x: *const i8) -> i32 {
6717    #[cfg(target_arch = "aarch64")]
6718    unsafe {
6719        return sdot32_i8(w, x);
6720    }
6721    #[cfg(target_arch = "x86_64")]
6722    unsafe {
6723        return i8dot32_avx2(w, x);
6724    }
6725    #[allow(unreachable_code)]
6726    unsafe {
6727        let mut s = 0i32;
6728        for k in 0..GROUP_SIZE {
6729            s += *w.add(k) as i32 * *x.add(k) as i32;
6730        }
6731        s
6732    }
6733}
6734
6735#[inline]
6736unsafe fn q1t_unpack_reg_u64s(codes: *const u8) -> (u64, u64, u64, u64) {
6737    let (s0, s1, s2, s3, s4, s5, s6) = unsafe {
6738        (
6739            SIGN5_U64[*codes as usize],
6740            SIGN5_U64[*codes.add(1) as usize],
6741            SIGN5_U64[*codes.add(2) as usize],
6742            SIGN5_U64[*codes.add(3) as usize],
6743            SIGN5_U64[*codes.add(4) as usize],
6744            SIGN5_U64[*codes.add(5) as usize],
6745            SIGN5_U64[*codes.add(6) as usize],
6746        )
6747    };
6748
6749    let u0 = s0 | (s1 << 40);
6750    let u1 = (s1 >> 24) | (s2 << 16) | (s3 << 56);
6751    let u2 = (s3 >> 8) | (s4 << 32);
6752    let u3 = (s4 >> 32) | (s5 << 8) | (s6 << 48);
6753
6754    (u0, u1, u2, u3)
6755}
6756
6757/// One q1t row's int8 base dot: `Σ_group s·dot(signs, xq)` (before the shared
6758/// `sx`). Direct register unpacking (zero stack stores/loads, no STLF stalls).
6759/// ARM SDOT.
6760#[cfg(target_arch = "aarch64")]
6761#[target_feature(enable = "neon,dotprod")]
6762unsafe fn q1t_dot_row_sdot(bytes: &[u8], r: usize, gpr: usize, xq: &[i8]) -> f32 {
6763    use core::arch::aarch64::*;
6764    use core::arch::asm;
6765    unsafe {
6766        const TILE: usize = cortiq_core::quant::Q1T_TILE;
6767        let mut acc = 0f32;
6768        let bytes_ptr = bytes.as_ptr();
6769        let xq_ptr = xq.as_ptr();
6770        let row_off = r * gpr * TILE;
6771
6772        let gpr2 = gpr & !1;
6773        let mut gi = 0;
6774        while gi < gpr2 {
6775            let off0 = row_off + gi * TILE;
6776            let off1 = off0 + TILE;
6777            let s0 = f16_to_f32(u16::from_le_bytes([
6778                *bytes_ptr.add(off0),
6779                *bytes_ptr.add(off0 + 1),
6780            ]));
6781            let s1 = f16_to_f32(u16::from_le_bytes([
6782                *bytes_ptr.add(off1),
6783                *bytes_ptr.add(off1 + 1),
6784            ]));
6785
6786            let (u0_0, u1_0, u2_0, u3_0) = q1t_unpack_reg_u64s(bytes_ptr.add(off0 + 2));
6787            let (u0_1, u1_1, u2_1, u3_1) = q1t_unpack_reg_u64s(bytes_ptr.add(off1 + 2));
6788
6789            let w0_0 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u0_0), vcreate_u64(u1_0)));
6790            let w1_0 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u2_0), vcreate_u64(u3_0)));
6791            let w0_1 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u0_1), vcreate_u64(u1_1)));
6792            let w1_1 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u2_1), vcreate_u64(u3_1)));
6793
6794            let x0_0 = vld1q_s8(xq_ptr.add(gi * GROUP_SIZE));
6795            let x1_0 = vld1q_s8(xq_ptr.add(gi * GROUP_SIZE + 16));
6796            let x0_1 = vld1q_s8(xq_ptr.add((gi + 1) * GROUP_SIZE));
6797            let x1_1 = vld1q_s8(xq_ptr.add((gi + 1) * GROUP_SIZE + 16));
6798
6799            let (mut a0_0, mut a1_0) = (vdupq_n_s32(0), vdupq_n_s32(0));
6800            let (mut a0_1, mut a1_1) = (vdupq_n_s32(0), vdupq_n_s32(0));
6801            asm!(
6802                "sdot {a0_0:v}.4s, {w0_0:v}.16b, {x0_0:v}.16b",
6803                "sdot {a1_0:v}.4s, {w1_0:v}.16b, {x1_0:v}.16b",
6804                "sdot {a0_1:v}.4s, {w0_1:v}.16b, {x0_1:v}.16b",
6805                "sdot {a1_1:v}.4s, {w1_1:v}.16b, {x1_1:v}.16b",
6806                a0_0 = inout(vreg) a0_0, a1_0 = inout(vreg) a1_0,
6807                a0_1 = inout(vreg) a0_1, a1_1 = inout(vreg) a1_1,
6808                w0_0 = in(vreg) w0_0, x0_0 = in(vreg) x0_0, w1_0 = in(vreg) w1_0, x1_0 = in(vreg) x1_0,
6809                w0_1 = in(vreg) w0_1, x0_1 = in(vreg) x0_1, w1_1 = in(vreg) w1_1, x1_1 = in(vreg) x1_1,
6810                options(pure, nomem, nostack),
6811            );
6812            let d0 = vaddvq_s32(vaddq_s32(a0_0, a1_0));
6813            let d1 = vaddvq_s32(vaddq_s32(a0_1, a1_1));
6814            acc += d0 as f32 * s0 + d1 as f32 * s1;
6815            gi += 2;
6816        }
6817
6818        if gi < gpr {
6819            let off = row_off + gi * TILE;
6820            let s = f16_to_f32(u16::from_le_bytes([
6821                *bytes_ptr.add(off),
6822                *bytes_ptr.add(off + 1),
6823            ]));
6824            let (u0, u1, u2, u3) = q1t_unpack_reg_u64s(bytes_ptr.add(off + 2));
6825            let w0 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u0), vcreate_u64(u1)));
6826            let w1 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u2), vcreate_u64(u3)));
6827            let x0 = vld1q_s8(xq_ptr.add(gi * GROUP_SIZE));
6828            let x1 = vld1q_s8(xq_ptr.add(gi * GROUP_SIZE + 16));
6829            let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
6830            asm!(
6831                "sdot {a0:v}.4s, {w0:v}.16b, {x0:v}.16b",
6832                "sdot {a1:v}.4s, {w1:v}.16b, {x1:v}.16b",
6833                a0 = inout(vreg) a0, a1 = inout(vreg) a1,
6834                w0 = in(vreg) w0, x0 = in(vreg) x0, w1 = in(vreg) w1, x1 = in(vreg) x1,
6835                options(pure, nomem, nostack),
6836            );
6837            let d = vaddvq_s32(vaddq_s32(a0, a1));
6838            acc += d as f32 * s;
6839        }
6840        acc
6841    }
6842}
6843
6844/// x86 AVX2 mirror of `q1t_dot_row_sdot` (maddubs int8 dot per group).
6845#[cfg(target_arch = "x86_64")]
6846#[target_feature(enable = "avx2")]
6847unsafe fn q1t_dot_row_avx2(bytes: &[u8], r: usize, gpr: usize, xq: &[i8]) -> f32 {
6848    use core::arch::x86_64::*;
6849    unsafe {
6850        const TILE: usize = cortiq_core::quant::Q1T_TILE;
6851        let mut acc = 0f32;
6852        let bytes_ptr = bytes.as_ptr();
6853        let xq_ptr = xq.as_ptr();
6854        let row_off = r * gpr * TILE;
6855
6856        let ones = _mm256_set1_epi16(1);
6857        for gi in 0..gpr {
6858            let off = row_off + gi * TILE;
6859            let s = f16_to_f32(u16::from_le_bytes([
6860                *bytes_ptr.add(off),
6861                *bytes_ptr.add(off + 1),
6862            ]));
6863            let (u0, u1, u2, u3) = q1t_unpack_reg_u64s(bytes_ptr.add(off + 2));
6864            let wv = _mm256_set_epi64x(u3 as i64, u2 as i64, u1 as i64, u0 as i64);
6865            let xv = _mm256_loadu_si256(xq_ptr.add(gi * GROUP_SIZE) as *const __m256i);
6866            let p16 = _mm256_maddubs_epi16(_mm256_abs_epi8(wv), _mm256_sign_epi8(xv, wv));
6867            let d256 = _mm256_madd_epi16(p16, ones);
6868            let d128 = _mm_add_epi32(
6869                _mm256_castsi256_si128(d256),
6870                _mm256_extracti128_si256(d256, 1),
6871            );
6872            let d64 = _mm_add_epi32(d128, _mm_shuffle_epi32(d128, 0xee));
6873            let d32 = _mm_cvtsi128_si32(_mm_add_epi32(d64, _mm_shuffle_epi32(d64, 0x55)));
6874            acc += d32 as f32 * s;
6875        }
6876        acc
6877    }
6878}
6879
6880/// VNNI twin of `q1t_dot_row_avx2` (see `dpbusd_hsum`).
6881#[cfg(target_arch = "x86_64")]
6882#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
6883unsafe fn q1t_dot_row_vnni(bytes: &[u8], r: usize, gpr: usize, xq: &[i8]) -> f32 {
6884    use core::arch::x86_64::*;
6885    // SAFETY: same tile/xq contracts as `q1t_dot_row_avx2`.
6886    unsafe {
6887        const TILE: usize = cortiq_core::quant::Q1T_TILE;
6888        let mut acc = 0f32;
6889        let bytes_ptr = bytes.as_ptr();
6890        let xq_ptr = xq.as_ptr();
6891        let row_off = r * gpr * TILE;
6892        for gi in 0..gpr {
6893            let off = row_off + gi * TILE;
6894            let s = f16_to_f32(u16::from_le_bytes([
6895                *bytes_ptr.add(off),
6896                *bytes_ptr.add(off + 1),
6897            ]));
6898            let (u0, u1, u2, u3) = q1t_unpack_reg_u64s(bytes_ptr.add(off + 2));
6899            let wv = _mm256_set_epi64x(u3 as i64, u2 as i64, u1 as i64, u0 as i64);
6900            let xv = _mm256_loadu_si256(xq_ptr.add(gi * GROUP_SIZE) as *const __m256i);
6901            let d = dpbusd_hsum(_mm256_abs_epi8(wv), _mm256_sign_epi8(xv, wv));
6902            acc += d as f32 * s;
6903        }
6904        acc
6905    }
6906}
6907
6908/// Per-row int8 base dot, dispatched once per row (matvec decode hot path).
6909/// Callers are gated by `a8w8_enabled()`, so the target-feature kernels are
6910/// reachable.
6911#[inline]
6912fn q1t_dot_row_i8(bytes: &[u8], r: usize, gpr: usize, xq: &[i8]) -> f32 {
6913    #[cfg(target_arch = "aarch64")]
6914    unsafe {
6915        return q1t_dot_row_sdot(bytes, r, gpr, xq);
6916    }
6917    #[cfg(target_arch = "x86_64")]
6918    unsafe {
6919        if vnni_tiles_enabled() {
6920            return q1t_dot_row_vnni(bytes, r, gpr, xq);
6921        }
6922        return q1t_dot_row_avx2(bytes, r, gpr, xq);
6923    }
6924    #[allow(unreachable_code)]
6925    {
6926        const TILE: usize = cortiq_core::quant::Q1T_TILE;
6927        let mut acc = 0f32;
6928        let mut sg = [0i8; GROUP_SIZE + 8]; // +8 slack for the u64-store unpack
6929        for gi in 0..gpr {
6930            let off = (r * gpr + gi) * TILE;
6931            let s = f16_to_f32(u16::from_le_bytes([bytes[off], bytes[off + 1]]));
6932            q1t_unpack_group_i8(bytes.as_ptr().wrapping_add(off + 2), &mut sg);
6933            let mut d = 0i32;
6934            for k in 0..GROUP_SIZE {
6935                d += sg[k] as i32 * xq[gi * GROUP_SIZE + k] as i32;
6936            }
6937            acc += d as f32 * s;
6938        }
6939        acc
6940    }
6941}
6942
6943/// Σ over a row's outliers of `value·x[col]` — the correction that adds the
6944/// overlay's exact weights on top of the base dot. INVARIANT: the encoder
6945/// writes ternary code 0 at every outlier position (`quantize_q1t`), so the
6946/// base contributes nothing there and this is a plain `value·x`, not
6947/// `(value − base)·x` — no scattered per-outlier scale read. Row `r`'s entries
6948/// are the contiguous slice `[row_ptr[r], row_ptr[r+1])`, so no binary search.
6949fn q1t_row_outlier_correction(
6950    bytes: &[u8],
6951    r: usize,
6952    rp_off: usize,
6953    entries_off: usize,
6954    has_ov: bool,
6955    x: &[f32],
6956) -> f32 {
6957    if !has_ov {
6958        return 0.0;
6959    }
6960    let (c0, c1) = (
6961        q1t_rowptr(bytes, rp_off, r),
6962        q1t_rowptr(bytes, rp_off, r + 1),
6963    );
6964    let mut corr = 0f32;
6965    for p in c0..c1 {
6966        let e = entries_off + p * 4;
6967        let col = u16::from_le_bytes([bytes[e], bytes[e + 1]]) as usize;
6968        let val = f16_to_f32(u16::from_le_bytes([bytes[e + 2], bytes[e + 3]]));
6969        corr += val * x[col];
6970    }
6971    corr
6972}
6973
6974/// Dequantize one q1t row into `buf[..cols]` via the sign LUT (no division),
6975/// then apply the row's outliers (its `[row_ptr[r], row_ptr[r+1])` slice).
6976/// Used by the batched (prefill) path where the decode amortizes over the batch.
6977fn q1t_dequant_row(
6978    bytes: &[u8],
6979    r: usize,
6980    gpr: usize,
6981    rp_off: usize,
6982    entries_off: usize,
6983    has_ov: bool,
6984    buf: &mut [f32],
6985) {
6986    const TILE: usize = cortiq_core::quant::Q1T_TILE;
6987    for g in 0..gpr {
6988        let off = (r * gpr + g) * TILE;
6989        let s = f16_to_f32(u16::from_le_bytes([bytes[off], bytes[off + 1]]));
6990        let codes = &bytes[off + 2..off + TILE];
6991        let bc = g * GROUP_SIZE;
6992        // 6 full bytes (30 codes) + a 7th byte holding the last 2.
6993        for bi in 0..6 {
6994            let lut = &SIGN5[codes[bi] as usize];
6995            let d = &mut buf[bc + bi * 5..bc + bi * 5 + 5];
6996            for i in 0..5 {
6997                d[i] = lut[i] * s;
6998            }
6999        }
7000        let lut = &SIGN5[codes[6] as usize];
7001        buf[bc + 30] = lut[0] * s;
7002        buf[bc + 31] = lut[1] * s;
7003    }
7004    if !has_ov {
7005        return;
7006    }
7007    let (c0, c1) = (
7008        q1t_rowptr(bytes, rp_off, r),
7009        q1t_rowptr(bytes, rp_off, r + 1),
7010    );
7011    for p in c0..c1 {
7012        let e = entries_off + p * 4;
7013        let col = u16::from_le_bytes([bytes[e], bytes[e + 1]]) as usize;
7014        buf[col] = f16_to_f32(u16::from_le_bytes([bytes[e + 2], bytes[e + 3]]));
7015    }
7016}
7017
7018/// Add the sparse outlier overlay onto a base dot already in `out` (the GPU
7019/// computes the ternary base; the overlay stays on the CPU — its entries are
7020/// few and its per-row gather doesn't vectorize on the GPU). Row-parallel.
7021fn q1t_add_overlay(
7022    bytes: &[u8],
7023    x: &[f32],
7024    rows: usize,
7025    cols: usize,
7026    out: &mut [f32],
7027    pool: Option<&Pool>,
7028) {
7029    const TILE: usize = cortiq_core::quant::Q1T_TILE;
7030    let gpr = cols / GROUP_SIZE;
7031    let (rp_off, ent_off, has_ov) = q1t_overlay(bytes, rows * gpr * TILE, rows);
7032    if !has_ov {
7033        return;
7034    }
7035    let out_addr = SendMut(out.as_mut_ptr());
7036    let run = move |start: usize, end: usize| {
7037        for r in start..end {
7038            let corr = q1t_row_outlier_correction(bytes, r, rp_off, ent_off, has_ov, x);
7039            // SAFETY: disjoint rows; add onto the base the GPU already wrote.
7040            unsafe { *out_addr.at(r) += corr };
7041        }
7042    };
7043    dispatch_rows(pool, rows, &run);
7044}
7045
7046/// Q1T row range via the A8W8 int8 path — shared activation split,
7047/// per-row: base SDOT dot + outlier correction + overlay.
7048#[allow(clippy::too_many_arguments)]
7049fn q1t_range_a8w8(
7050    bytes: &[u8],
7051    gpr: usize,
7052    rp_off: usize,
7053    ent_off: usize,
7054    has_ov: bool,
7055    act: &SplitAct,
7056    x: &[f32],
7057    out: SendMut,
7058    start: usize,
7059    end: usize,
7060) {
7061    for r in start..end {
7062        let mut acc = q1t_dot_row_i8(bytes, r, gpr, &act.xq) * act.sx;
7063        for &(j, xv) in &act.outliers {
7064            acc += q1t_base_weight(bytes, r, gpr, j) * xv;
7065        }
7066        acc += q1t_row_outlier_correction(bytes, r, rp_off, ent_off, has_ov, x);
7067        // SAFETY: disjoint row ranges per worker.
7068        unsafe { *out.at(r) = acc };
7069    }
7070}
7071
7072/// Q1T row range via the f32 path (no SDOT) — for matvec_many batched
7073/// dispatch when a8w8 is unavailable.
7074#[allow(clippy::too_many_arguments)]
7075fn q1t_range_f32_batch(
7076    bytes: &[u8],
7077    gpr: usize,
7078    rp_off: usize,
7079    ent_off: usize,
7080    has_ov: bool,
7081    x: &[f32],
7082    out: SendMut,
7083    start: usize,
7084    end: usize,
7085) {
7086    const TILE: usize = cortiq_core::quant::Q1T_TILE;
7087    let mut sg = [0f32; GROUP_SIZE];
7088    for r in start..end {
7089        let mut acc = 0f32;
7090        for g in 0..gpr {
7091            let off = (r * gpr + g) * TILE;
7092            let s = f16_to_f32(u16::from_le_bytes([bytes[off], bytes[off + 1]]));
7093            let codes = &bytes[off + 2..off + TILE];
7094            let xg = &x[g * GROUP_SIZE..g * GROUP_SIZE + GROUP_SIZE];
7095            for bi in 0..6 {
7096                sg[bi * 5..bi * 5 + 5].copy_from_slice(&SIGN5[codes[bi] as usize]);
7097            }
7098            let lut = &SIGN5[codes[6] as usize];
7099            sg[30] = lut[0];
7100            sg[31] = lut[1];
7101            let mut gsum = 0f32;
7102            for k in 0..GROUP_SIZE {
7103                gsum += sg[k] * xg[k];
7104            }
7105            acc += s * gsum;
7106        }
7107        acc += q1t_row_outlier_correction(bytes, r, rp_off, ent_off, has_ov, x);
7108        // SAFETY: disjoint row ranges per worker.
7109        unsafe { *out.at(r) = acc };
7110    }
7111}
7112
7113/// Ternary (q1t) matvec — decode+dot straight from mmap, one group at a time:
7114/// no per-ROW buffer, no division (the sign LUT), and a tiny per-group sign
7115/// buffer so the 32-wide dot vectorizes. This is the decode hot path.
7116fn q1t_matvec(
7117    bytes: &[u8],
7118    x: &[f32],
7119    rows: usize,
7120    cols: usize,
7121    out: &mut [f32],
7122    pool: Option<&Pool>,
7123) {
7124    debug_assert_eq!(out.len(), rows);
7125    const TILE: usize = cortiq_core::quant::Q1T_TILE;
7126    let gpr = cols / GROUP_SIZE;
7127    let (rp_off, ent_off, has_ov) = q1t_overlay(bytes, rows * gpr * TILE, rows);
7128    let out_addr = SendMut(out.as_mut_ptr());
7129    // int8 SDOT base dot (ARM dotprod): ~4× the f32 arithmetic. x → i8 once
7130    // (`split_act`), activation outliers added back exactly in f32, weight
7131    // overlay on top. ARM SDOT / x86 AVX2; CMF_SDOT=0 keeps the exact f32 path.
7132    if a8w8_enabled() {
7133        let act = split_act(x);
7134        let act = &act;
7135        let run = move |start: usize, end: usize| {
7136            for r in start..end {
7137                let mut acc = q1t_dot_row_i8(bytes, r, gpr, &act.xq) * act.sx;
7138                for &(j, xv) in &act.outliers {
7139                    acc += q1t_base_weight(bytes, r, gpr, j) * xv;
7140                }
7141                acc += q1t_row_outlier_correction(bytes, r, rp_off, ent_off, has_ov, x);
7142                // SAFETY: disjoint row ranges per worker.
7143                unsafe { *out_addr.at(r) = acc };
7144            }
7145        };
7146        dispatch_rows(pool, rows, &run);
7147        return;
7148    }
7149    let run = move |start: usize, end: usize| {
7150        // Per-group signs, unpacked contiguously so the dot below is a clean
7151        // 32-wide reduction the autovectorizer turns into f32x4 FMAs — the
7152        // 5-values-per-byte base-3 layout won't SIMD in place.
7153        let mut sg = [0f32; GROUP_SIZE];
7154        for r in start..end {
7155            let mut acc = 0f32;
7156            for g in 0..gpr {
7157                let off = (r * gpr + g) * TILE;
7158                let s = f16_to_f32(u16::from_le_bytes([bytes[off], bytes[off + 1]]));
7159                let codes = &bytes[off + 2..off + TILE];
7160                let xg = &x[g * GROUP_SIZE..g * GROUP_SIZE + GROUP_SIZE];
7161                for bi in 0..6 {
7162                    sg[bi * 5..bi * 5 + 5].copy_from_slice(&SIGN5[codes[bi] as usize]);
7163                }
7164                let lut = &SIGN5[codes[6] as usize];
7165                sg[30] = lut[0];
7166                sg[31] = lut[1];
7167                let mut gsum = 0f32;
7168                for k in 0..GROUP_SIZE {
7169                    gsum += sg[k] * xg[k];
7170                }
7171                acc += s * gsum;
7172            }
7173            acc += q1t_row_outlier_correction(bytes, r, rp_off, ent_off, has_ov, x);
7174            unsafe { *out_addr.at(r) = acc };
7175        }
7176    };
7177    dispatch_rows(pool, rows, &run);
7178}
7179
7180/// Fused-pair twin of `q1t_dot_row_sdot`: ONE register unpack of the
7181/// ternary codes serves BOTH activation streams (the unpack chain is
7182/// the dominant per-row cost — MTP verify pairs paid it twice). Per
7183/// stream the group order and f32 accumulation match the single-row
7184/// kernel exactly, so pair == 2×matvec bit-for-bit.
7185#[cfg(target_arch = "aarch64")]
7186#[target_feature(enable = "neon,dotprod")]
7187unsafe fn q1t_dot_row_sdot2(bytes: &[u8], r: usize, gpr: usize, xa: &[i8], xb: &[i8]) -> [f32; 2] {
7188    use core::arch::aarch64::*;
7189    use core::arch::asm;
7190    // SAFETY: same slice-length contracts as `q1t_dot_row_sdot`, ×2.
7191    unsafe {
7192        const TILE: usize = cortiq_core::quant::Q1T_TILE;
7193        let bytes_ptr = bytes.as_ptr();
7194        let row_off = r * gpr * TILE;
7195        let xp = [xa.as_ptr(), xb.as_ptr()];
7196        let mut acc = [0f32; 2];
7197        macro_rules! sdot2 {
7198            ($w0:expr, $w1:expr, $x:expr) => {{
7199                let x0 = vld1q_s8($x);
7200                let x1 = vld1q_s8($x.add(16));
7201                let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
7202                asm!(
7203                    "sdot {a0:v}.4s, {w0:v}.16b, {x0:v}.16b",
7204                    "sdot {a1:v}.4s, {w1:v}.16b, {x1:v}.16b",
7205                    a0 = inout(vreg) a0, a1 = inout(vreg) a1,
7206                    w0 = in(vreg) $w0, x0 = in(vreg) x0, w1 = in(vreg) $w1, x1 = in(vreg) x1,
7207                    options(pure, nomem, nostack),
7208                );
7209                vaddvq_s32(vaddq_s32(a0, a1))
7210            }};
7211        }
7212        let gpr2 = gpr & !1;
7213        let mut gi = 0;
7214        while gi < gpr2 {
7215            let off0 = row_off + gi * TILE;
7216            let off1 = off0 + TILE;
7217            let s0 = f16_to_f32(u16::from_le_bytes([
7218                *bytes_ptr.add(off0),
7219                *bytes_ptr.add(off0 + 1),
7220            ]));
7221            let s1 = f16_to_f32(u16::from_le_bytes([
7222                *bytes_ptr.add(off1),
7223                *bytes_ptr.add(off1 + 1),
7224            ]));
7225            let (u0_0, u1_0, u2_0, u3_0) = q1t_unpack_reg_u64s(bytes_ptr.add(off0 + 2));
7226            let (u0_1, u1_1, u2_1, u3_1) = q1t_unpack_reg_u64s(bytes_ptr.add(off1 + 2));
7227            let w0_0 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u0_0), vcreate_u64(u1_0)));
7228            let w1_0 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u2_0), vcreate_u64(u3_0)));
7229            let w0_1 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u0_1), vcreate_u64(u1_1)));
7230            let w1_1 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u2_1), vcreate_u64(u3_1)));
7231            for k in 0..2 {
7232                let d0 = sdot2!(w0_0, w1_0, xp[k].add(gi * GROUP_SIZE));
7233                let d1 = sdot2!(w0_1, w1_1, xp[k].add((gi + 1) * GROUP_SIZE));
7234                acc[k] += d0 as f32 * s0 + d1 as f32 * s1;
7235            }
7236            gi += 2;
7237        }
7238        if gi < gpr {
7239            let off = row_off + gi * TILE;
7240            let s = f16_to_f32(u16::from_le_bytes([
7241                *bytes_ptr.add(off),
7242                *bytes_ptr.add(off + 1),
7243            ]));
7244            let (u0, u1, u2, u3) = q1t_unpack_reg_u64s(bytes_ptr.add(off + 2));
7245            let w0 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u0), vcreate_u64(u1)));
7246            let w1 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u2), vcreate_u64(u3)));
7247            for k in 0..2 {
7248                let d = sdot2!(w0, w1, xp[k].add(gi * GROUP_SIZE));
7249                acc[k] += d as f32 * s;
7250            }
7251        }
7252        acc
7253    }
7254}
7255
7256/// Fused Q1T pair matvec: ONE pass over the rows serves both
7257/// activation streams — on ARM the ternary register unpack happens
7258/// once per tile pair (`q1t_dot_row_sdot2`); elsewhere the second dot
7259/// rides the row's L1-warm tile bytes. Per stream the math matches
7260/// `q1t_matvec` exactly.
7261fn q1t_matvec2(
7262    bytes: &[u8],
7263    x1: &[f32],
7264    x2: &[f32],
7265    rows: usize,
7266    cols: usize,
7267    o1: &mut [f32],
7268    o2: &mut [f32],
7269    pool: Option<&Pool>,
7270) {
7271    debug_assert_eq!(o1.len(), rows);
7272    debug_assert_eq!(o2.len(), rows);
7273    const TILE: usize = cortiq_core::quant::Q1T_TILE;
7274    let gpr = cols / GROUP_SIZE;
7275    let (rp_off, ent_off, has_ov) = q1t_overlay(bytes, rows * gpr * TILE, rows);
7276    let out1 = SendMut(o1.as_mut_ptr());
7277    let out2 = SendMut(o2.as_mut_ptr());
7278    if a8w8_enabled() {
7279        let a1 = split_act(x1);
7280        let a2 = split_act(x2);
7281        let (a1, a2) = (&a1, &a2);
7282        let run = move |start: usize, end: usize| {
7283            for r in start..end {
7284                #[cfg(target_arch = "aarch64")]
7285                // a8w8 on aarch64 ⇔ sdot_enabled(), so the kernel's
7286                // target features are present.
7287                let ds = unsafe { q1t_dot_row_sdot2(bytes, r, gpr, &a1.xq, &a2.xq) };
7288                #[cfg(not(target_arch = "aarch64"))]
7289                let ds = [
7290                    q1t_dot_row_i8(bytes, r, gpr, &a1.xq),
7291                    q1t_dot_row_i8(bytes, r, gpr, &a2.xq),
7292                ];
7293                let mut acc1 = ds[0] * a1.sx;
7294                for &(j, xv) in &a1.outliers {
7295                    acc1 += q1t_base_weight(bytes, r, gpr, j) * xv;
7296                }
7297                acc1 += q1t_row_outlier_correction(bytes, r, rp_off, ent_off, has_ov, x1);
7298                let mut acc2 = ds[1] * a2.sx;
7299                for &(j, xv) in &a2.outliers {
7300                    acc2 += q1t_base_weight(bytes, r, gpr, j) * xv;
7301                }
7302                acc2 += q1t_row_outlier_correction(bytes, r, rp_off, ent_off, has_ov, x2);
7303                // SAFETY: disjoint row ranges per worker.
7304                unsafe {
7305                    *out1.at(r) = acc1;
7306                    *out2.at(r) = acc2;
7307                }
7308            }
7309        };
7310        dispatch_rows(pool, rows, &run);
7311        return;
7312    }
7313    let run = move |start: usize, end: usize| {
7314        // Exact path (CMF_SDOT=0): unpack the sign LUT once per group,
7315        // dot both streams — same op order per stream as `q1t_matvec`.
7316        let mut sg = [0f32; GROUP_SIZE];
7317        for r in start..end {
7318            let mut acc1 = 0f32;
7319            let mut acc2 = 0f32;
7320            for g in 0..gpr {
7321                let off = (r * gpr + g) * TILE;
7322                let s = f16_to_f32(u16::from_le_bytes([bytes[off], bytes[off + 1]]));
7323                let codes = &bytes[off + 2..off + TILE];
7324                for bi in 0..6 {
7325                    sg[bi * 5..bi * 5 + 5].copy_from_slice(&SIGN5[codes[bi] as usize]);
7326                }
7327                let lut = &SIGN5[codes[6] as usize];
7328                sg[30] = lut[0];
7329                sg[31] = lut[1];
7330                let xg1 = &x1[g * GROUP_SIZE..g * GROUP_SIZE + GROUP_SIZE];
7331                let xg2 = &x2[g * GROUP_SIZE..g * GROUP_SIZE + GROUP_SIZE];
7332                let mut gsum1 = 0f32;
7333                for k in 0..GROUP_SIZE {
7334                    gsum1 += sg[k] * xg1[k];
7335                }
7336                acc1 += s * gsum1;
7337                let mut gsum2 = 0f32;
7338                for k in 0..GROUP_SIZE {
7339                    gsum2 += sg[k] * xg2[k];
7340                }
7341                acc2 += s * gsum2;
7342            }
7343            acc1 += q1t_row_outlier_correction(bytes, r, rp_off, ent_off, has_ov, x1);
7344            acc2 += q1t_row_outlier_correction(bytes, r, rp_off, ent_off, has_ov, x2);
7345            // SAFETY: disjoint row ranges per worker.
7346            unsafe {
7347                *out1.at(r) = acc1;
7348                *out2.at(r) = acc2;
7349            }
7350        }
7351    };
7352    dispatch_rows(pool, rows, &run);
7353}
7354
7355/// Ternary (q1t) matmat (prefill) — dequant each row once, dot the whole
7356/// batch against it (amortizes the per-row decode).
7357fn q1t_matmat(
7358    bytes: &[u8],
7359    xs: &[f32],
7360    b: usize,
7361    rows: usize,
7362    cols: usize,
7363    out: &mut [f32],
7364    pool: Option<&Pool>,
7365) {
7366    debug_assert_eq!(out.len(), b * rows);
7367    const TILE: usize = cortiq_core::quant::Q1T_TILE;
7368    let gpr = cols / GROUP_SIZE;
7369    let (rp_off, ent_off, has_ov) = q1t_overlay(bytes, rows * gpr * TILE, rows);
7370    let out_addr = SendMut(out.as_mut_ptr());
7371    // int8 prefill (ARM SDOT / x86 AVX2): quantize the B inputs once, unpack
7372    // each weight row's signs to i8 ONCE, then int8-dot against every input —
7373    // the row sign-decode amortizes over the whole batch. CMF_SDOT=0 → f32.
7374    if a8w8_enabled() {
7375        let acts: Vec<SplitAct> = (0..b)
7376            .map(|bi| split_act(&xs[bi * cols..(bi + 1) * cols]))
7377            .collect();
7378        let acts = &acts;
7379        let run = move |start: usize, end: usize| {
7380            let mut sg = vec![0i8; cols + 8]; // row signs, i8 (+8 unpack slack)
7381            let mut sc = vec![0f32; gpr]; // per-group scales
7382            let mut accs = vec![0f32; b]; // per-batch accumulators, reused per row
7383            for r in start..end {
7384                for g in 0..gpr {
7385                    let off = (r * gpr + g) * TILE;
7386                    sc[g] = f16_to_f32(u16::from_le_bytes([bytes[off], bytes[off + 1]]));
7387                    q1t_unpack_group_i8(
7388                        bytes.as_ptr().wrapping_add(off + 2),
7389                        &mut sg[g * GROUP_SIZE..],
7390                    );
7391                }
7392                for bi in 0..b {
7393                    let act = &acts[bi];
7394                    let mut isum = 0f32;
7395                    for g in 0..gpr {
7396                        let d = q1t_i8dot32(
7397                            sg.as_ptr().wrapping_add(g * GROUP_SIZE),
7398                            act.xq.as_ptr().wrapping_add(g * GROUP_SIZE),
7399                        );
7400                        isum += d as f32 * sc[g];
7401                    }
7402                    let mut acc = isum * act.sx;
7403                    for &(j, xv) in &act.outliers {
7404                        acc += q1t_base_weight(bytes, r, gpr, j) * xv;
7405                    }
7406                    accs[bi] = acc;
7407                }
7408                // Overlay ONCE per row for the whole batch: read each (col, val)
7409                // from mmap a single time (was b× — the re-read dominated prefill)
7410                // and fan it out over the batch via the cached inputs.
7411                if has_ov {
7412                    let (c0, c1) = (
7413                        q1t_rowptr(bytes, rp_off, r),
7414                        q1t_rowptr(bytes, rp_off, r + 1),
7415                    );
7416                    for p in c0..c1 {
7417                        let e = ent_off + p * 4;
7418                        let col = u16::from_le_bytes([bytes[e], bytes[e + 1]]) as usize;
7419                        let val = f16_to_f32(u16::from_le_bytes([bytes[e + 2], bytes[e + 3]]));
7420                        for bi in 0..b {
7421                            accs[bi] += val * xs[bi * cols + col];
7422                        }
7423                    }
7424                }
7425                for bi in 0..b {
7426                    unsafe { *out_addr.at(bi * rows + r) = accs[bi] };
7427                }
7428            }
7429        };
7430        dispatch_rows(pool, rows, &run);
7431        return;
7432    }
7433    let run = move |start: usize, end: usize| {
7434        let mut buf = vec![0f32; cols];
7435        for r in start..end {
7436            q1t_dequant_row(bytes, r, gpr, rp_off, ent_off, has_ov, &mut buf);
7437            for bi in 0..b {
7438                let xr = &xs[bi * cols..(bi + 1) * cols];
7439                let mut acc = 0f32;
7440                for j in 0..cols {
7441                    acc += buf[j] * xr[j];
7442                }
7443                unsafe { *out_addr.at(bi * rows + r) = acc };
7444            }
7445        }
7446    };
7447    dispatch_rows(pool, rows, &run);
7448}
7449
7450fn q1_matvec(
7451    bytes: &[u8],
7452    x: &[f32],
7453    rows: usize,
7454    cols: usize,
7455    out: &mut [f32],
7456    pool: Option<&Pool>,
7457) {
7458    debug_assert_eq!(out.len(), rows);
7459    let gpr = cols / GROUP_SIZE;
7460    let out_addr = SendMut(out.as_mut_ptr());
7461    if a8w8_enabled() {
7462        let act = split_act(x);
7463        let gsum = q1_group_sums(&act.xq, gpr);
7464        let (act, gsum) = (&act, &gsum);
7465        let run = move |start: usize, end: usize| {
7466            q1_range_a8w8(bytes, gpr, act, gsum, out_addr, start, end)
7467        };
7468        dispatch_rows(pool, rows, &run);
7469        return;
7470    }
7471    let run = move |start: usize, end: usize| q1_range_f32(bytes, gpr, x, out_addr, start, end);
7472    dispatch_rows(pool, rows, &run);
7473}
7474
7475/// Fused two-input q1 matvec (weights read once per pair).
7476#[allow(clippy::too_many_arguments)]
7477fn q1_matvec2(
7478    bytes: &[u8],
7479    x1: &[f32],
7480    x2: &[f32],
7481    rows: usize,
7482    cols: usize,
7483    o1: &mut [f32],
7484    o2: &mut [f32],
7485    pool: Option<&Pool>,
7486) {
7487    let gpr = cols / GROUP_SIZE;
7488    let p1 = SendMut(o1.as_mut_ptr());
7489    let p2 = SendMut(o2.as_mut_ptr());
7490    if a8w8_enabled() {
7491        let a1 = split_act(x1);
7492        let a2 = split_act(x2);
7493        let g1 = q1_group_sums(&a1.xq, gpr);
7494        let g2 = q1_group_sums(&a2.xq, gpr);
7495        let (a1, a2, g1, g2) = (&a1, &a2, &g1, &g2);
7496        let run = move |start: usize, end: usize| {
7497            for r in start..end {
7498                let mut v1 = dot_q1_row_i8(bytes, r, gpr, &a1.xq, g1) * a1.sx;
7499                let mut v2 = dot_q1_row_i8(bytes, r, gpr, &a2.xq, g2) * a2.sx;
7500                for &(j, xv) in &a1.outliers {
7501                    let (w, s) = q1_outlier(bytes, r, gpr, j);
7502                    v1 += w * s * xv;
7503                }
7504                for &(j, xv) in &a2.outliers {
7505                    let (w, s) = q1_outlier(bytes, r, gpr, j);
7506                    v2 += w * s * xv;
7507                }
7508                // SAFETY: disjoint row ranges per worker.
7509                unsafe {
7510                    *p1.at(r) = v1;
7511                    *p2.at(r) = v2;
7512                }
7513            }
7514        };
7515        dispatch_rows(pool, rows, &run);
7516        return;
7517    }
7518    let run = move |start: usize, end: usize| {
7519        for r in start..end {
7520            // SAFETY: disjoint row ranges per worker.
7521            unsafe {
7522                *p1.at(r) = q1_row_exact(bytes, r, gpr, x1);
7523                *p2.at(r) = q1_row_exact(bytes, r, gpr, x2);
7524            }
7525        }
7526    };
7527    dispatch_rows(pool, rows, &run);
7528}
7529
7530/// Batched q1 matmat: each row's tiles stream once per microbatch.
7531#[allow(clippy::too_many_arguments)]
7532fn q1_matmat(
7533    bytes: &[u8],
7534    xs_all: &[f32],
7535    b: usize,
7536    rows: usize,
7537    cols: usize,
7538    out: &mut [f32],
7539    pool: Option<&Pool>,
7540) {
7541    debug_assert_eq!(out.len(), b * rows);
7542    let gpr = cols / GROUP_SIZE;
7543    let out_addr = SendMut(out.as_mut_ptr());
7544    if a8w8_enabled() {
7545        let acts: Vec<(SplitAct, Vec<i32>)> = (0..b)
7546            .map(|bi| {
7547                let act = split_act(&xs_all[bi * cols..(bi + 1) * cols]);
7548                let gsum = q1_group_sums(&act.xq, gpr);
7549                (act, gsum)
7550            })
7551            .collect();
7552        let acts = &acts;
7553        #[cfg(target_arch = "x86_64")]
7554        let blocked_ok = avx2_enabled() && blocked_enabled();
7555        #[cfg(target_arch = "aarch64")]
7556        let blocked_ok = sdot_enabled() && blocked_enabled();
7557        let run = move |start: usize, end: usize| {
7558            for r in start..end {
7559                let mut bi = 0usize;
7560                // Blocked 1×4: the unpacked bit mask serves four
7561                // activation streams per group.
7562                #[cfg(target_arch = "aarch64")]
7563                if blocked_ok {
7564                    while bi + 4 <= acts.len() {
7565                        let xs = [
7566                            acts[bi].0.xq.as_slice(),
7567                            acts[bi + 1].0.xq.as_slice(),
7568                            acts[bi + 2].0.xq.as_slice(),
7569                            acts[bi + 3].0.xq.as_slice(),
7570                        ];
7571                        let gs = [
7572                            acts[bi].1.as_slice(),
7573                            acts[bi + 1].1.as_slice(),
7574                            acts[bi + 2].1.as_slice(),
7575                            acts[bi + 3].1.as_slice(),
7576                        ];
7577                        let d = unsafe { dot_q1_row_1x4_sdot(bytes, r, gpr, xs, gs) };
7578                        for k in 0..4 {
7579                            let (act, _) = &acts[bi + k];
7580                            let mut acc = d[k] * act.sx;
7581                            for &(j, xv) in &act.outliers {
7582                                let (w, sc) = q1_outlier(bytes, r, gpr, j);
7583                                acc += w * sc * xv;
7584                            }
7585                            // SAFETY: disjoint (bi, r) cells per worker.
7586                            unsafe { *out_addr.at((bi + k) * rows + r) = acc };
7587                        }
7588                        bi += 4;
7589                    }
7590                }
7591                #[cfg(target_arch = "x86_64")]
7592                if blocked_ok {
7593                    while bi + 4 <= acts.len() {
7594                        let xs = [
7595                            acts[bi].0.xq.as_slice(),
7596                            acts[bi + 1].0.xq.as_slice(),
7597                            acts[bi + 2].0.xq.as_slice(),
7598                            acts[bi + 3].0.xq.as_slice(),
7599                        ];
7600                        let gs = [
7601                            acts[bi].1.as_slice(),
7602                            acts[bi + 1].1.as_slice(),
7603                            acts[bi + 2].1.as_slice(),
7604                            acts[bi + 3].1.as_slice(),
7605                        ];
7606                        let d = unsafe {
7607                            if vnni_tiles_enabled() {
7608                                dot_q1_row_1x4_vnni(bytes, r, gpr, xs, gs)
7609                            } else {
7610                                dot_q1_row_1x4_avx2(bytes, r, gpr, xs, gs)
7611                            }
7612                        };
7613                        for k in 0..4 {
7614                            let (act, _) = &acts[bi + k];
7615                            let mut acc = d[k] * act.sx;
7616                            for &(j, xv) in &act.outliers {
7617                                let (w, sc) = q1_outlier(bytes, r, gpr, j);
7618                                acc += w * sc * xv;
7619                            }
7620                            // SAFETY: disjoint (bi, r) cells per worker.
7621                            unsafe { *out_addr.at((bi + k) * rows + r) = acc };
7622                        }
7623                        bi += 4;
7624                    }
7625                }
7626                while bi < acts.len() {
7627                    let (act, gsum) = &acts[bi];
7628                    let mut acc = dot_q1_row_i8(bytes, r, gpr, &act.xq, gsum) * act.sx;
7629                    for &(j, xv) in &act.outliers {
7630                        let (w, s) = q1_outlier(bytes, r, gpr, j);
7631                        acc += w * s * xv;
7632                    }
7633                    // SAFETY: disjoint (bi, r) cells per worker range.
7634                    unsafe { *out_addr.at(bi * rows + r) = acc };
7635                    bi += 1;
7636                }
7637            }
7638        };
7639        dispatch_rows(pool, rows, &run);
7640        return;
7641    }
7642    let run = move |start: usize, end: usize| {
7643        for r in start..end {
7644            for bi in 0..b {
7645                let x = &xs_all[bi * cols..(bi + 1) * cols];
7646                // SAFETY: disjoint (bi, r) cells per worker range.
7647                unsafe { *out_addr.at(bi * rows + r) = q1_row_exact(bytes, r, gpr, x) };
7648            }
7649        }
7650    };
7651    dispatch_rows(pool, rows, &run);
7652}
7653
7654/// Fused q4_block matvec straight from the mapped bytes. SDOT path when
7655/// dotprod is available (port of vmfcore `dot_q4_block_sdot`, measured
7656/// +23% on q4 decode): nibbles → centered i8, int8×int8 `sdot` per
7657/// 32-group, exact outlier correction — the same A8W8 contract as q8.
7658/// `CMF_SDOT=0` keeps the exact scalar path.
7659fn q4matvec(
7660    bytes: &[u8],
7661    x: &[f32],
7662    rows: usize,
7663    cols: usize,
7664    out: &mut [f32],
7665    pool: Option<&Pool>,
7666) {
7667    debug_assert_eq!(out.len(), rows);
7668    let (packed, scales) = q4_split(bytes, rows, cols);
7669    let gpr = cols / GROUP_SIZE;
7670    let out_addr = SendMut(out.as_mut_ptr());
7671
7672    if a8w8_enabled() {
7673        let act = split_act(x);
7674        let run = move |start: usize, end: usize| {
7675            q4_range_a8w8(packed, scales, gpr, cols, &act, out_addr, start, end)
7676        };
7677        dispatch_rows(pool, rows, &run);
7678        return;
7679    }
7680
7681    let run =
7682        move |start: usize, end: usize| q4_range_f32(packed, scales, gpr, x, out_addr, start, end);
7683    dispatch_rows(pool, rows, &run);
7684}
7685
7686/// One q4 row via the A8W8 int8 path — SDOT on ARM, AVX2 maddubs on
7687/// x86 (scalar fallback is unreachable: callers gate on a8w8_enabled).
7688#[inline]
7689#[allow(unreachable_code)]
7690/// One UNPACKED q4 row (centered i8 in `buf`) against four activation
7691/// streams: the 32-byte weight chunk and its abs() load once per group,
7692/// the per-group f16 scale decodes once — four maddubs+reduce chains
7693/// instead of four full (load, abs, dot) rounds.
7694#[cfg(target_arch = "x86_64")]
7695#[target_feature(enable = "avx2")]
7696unsafe fn dot_q4b_row_1x4_avx2(
7697    buf: &[u8],
7698    scales: &[u8],
7699    g0: usize,
7700    gpr: usize,
7701    xs: [&[i8]; 4],
7702) -> [f32; 4] {
7703    // SAFETY: callers uphold buffer contracts (buf.len() == gpr·32).
7704    unsafe {
7705        use core::arch::x86_64::*;
7706        let ones = _mm256_set1_epi16(1);
7707        let mut acc = [0f32; 4];
7708        for gi in 0..gpr {
7709            let s = f16_to_f32(u16::from_le_bytes([
7710                scales[(g0 + gi) * 2],
7711                scales[(g0 + gi) * 2 + 1],
7712            ]));
7713            let w = _mm256_loadu_si256(buf.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
7714            let aw = _mm256_abs_epi8(w);
7715            for (k, xq) in xs.iter().enumerate() {
7716                let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
7717                let p16 = _mm256_maddubs_epi16(aw, _mm256_sign_epi8(x, w));
7718                let d = _mm256_madd_epi16(p16, ones);
7719                let hi128 = _mm256_extracti128_si256::<1>(d);
7720                let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
7721                let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
7722                let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
7723                acc[k] += _mm_cvtsi128_si32(s32) as f32 * s;
7724            }
7725        }
7726        acc
7727    }
7728}
7729
7730/// VNNI twin of `dot_q4b_row_1x4_avx2` (see `dpbusd_hsum`).
7731#[cfg(target_arch = "x86_64")]
7732#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
7733unsafe fn dot_q4b_row_1x4_vnni(
7734    buf: &[u8],
7735    scales: &[u8],
7736    g0: usize,
7737    gpr: usize,
7738    xs: [&[i8]; 4],
7739) -> [f32; 4] {
7740    // SAFETY: callers uphold buffer contracts (buf.len() == gpr·32).
7741    unsafe {
7742        use core::arch::x86_64::*;
7743        let mut acc = [0f32; 4];
7744        for gi in 0..gpr {
7745            let s = f16_to_f32(u16::from_le_bytes([
7746                scales[(g0 + gi) * 2],
7747                scales[(g0 + gi) * 2 + 1],
7748            ]));
7749            let w = _mm256_loadu_si256(buf.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
7750            let aw = _mm256_abs_epi8(w);
7751            for (k, xq) in xs.iter().enumerate() {
7752                let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
7753                let d = dpbusd_hsum(aw, _mm256_sign_epi8(x, w));
7754                acc[k] += d as f32 * s;
7755            }
7756        }
7757        acc
7758    }
7759}
7760
7761/// The vbit flavor of the blocked 1×4: the per-activation A8W8 scale
7762/// folds in PER GROUP as `(d·sx)·s` — bit-matching the single-matvec
7763/// accumulation order (the q4_block flavor applies sx once at the end,
7764/// matching ITS single path; the two conventions are historical and
7765/// each blocked leg must mirror its own).
7766#[cfg(target_arch = "x86_64")]
7767#[target_feature(enable = "avx2")]
7768unsafe fn dot_q4b_row_1x4_sx_avx2(
7769    buf: &[u8],
7770    scales: &[u8],
7771    g0: usize,
7772    gpr: usize,
7773    xs: [&[i8]; 4],
7774    sxs: [f32; 4],
7775) -> [f32; 4] {
7776    // SAFETY: callers uphold buffer contracts (buf.len() == gpr·32).
7777    unsafe {
7778        use core::arch::x86_64::*;
7779        let ones = _mm256_set1_epi16(1);
7780        let mut acc = [0f32; 4];
7781        for gi in 0..gpr {
7782            let s = f16_to_f32(u16::from_le_bytes([
7783                scales[(g0 + gi) * 2],
7784                scales[(g0 + gi) * 2 + 1],
7785            ]));
7786            let w = _mm256_loadu_si256(buf.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
7787            let aw = _mm256_abs_epi8(w);
7788            for (k, xq) in xs.iter().enumerate() {
7789                let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
7790                let p16 = _mm256_maddubs_epi16(aw, _mm256_sign_epi8(x, w));
7791                let d = _mm256_madd_epi16(p16, ones);
7792                let hi128 = _mm256_extracti128_si256::<1>(d);
7793                let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
7794                let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
7795                let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
7796                acc[k] += (_mm_cvtsi128_si32(s32) as f32 * sxs[k]) * s;
7797            }
7798        }
7799        acc
7800    }
7801}
7802
7803/// VNNI twin of `dot_q4b_row_1x4_sx_avx2` (see `dpbusd_hsum`; the
7804/// per-group `(d·sx)·s` fold mirrors the vbit single path).
7805#[cfg(target_arch = "x86_64")]
7806#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
7807unsafe fn dot_q4b_row_1x4_sx_vnni(
7808    buf: &[u8],
7809    scales: &[u8],
7810    g0: usize,
7811    gpr: usize,
7812    xs: [&[i8]; 4],
7813    sxs: [f32; 4],
7814) -> [f32; 4] {
7815    // SAFETY: callers uphold buffer contracts (buf.len() == gpr·32).
7816    unsafe {
7817        use core::arch::x86_64::*;
7818        let mut acc = [0f32; 4];
7819        for gi in 0..gpr {
7820            let s = f16_to_f32(u16::from_le_bytes([
7821                scales[(g0 + gi) * 2],
7822                scales[(g0 + gi) * 2 + 1],
7823            ]));
7824            let w = _mm256_loadu_si256(buf.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
7825            let aw = _mm256_abs_epi8(w);
7826            for (k, xq) in xs.iter().enumerate() {
7827                let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
7828                let d = dpbusd_hsum(aw, _mm256_sign_epi8(x, w));
7829                acc[k] += (d as f32 * sxs[k]) * s;
7830            }
7831        }
7832        acc
7833    }
7834}
7835
7836#[allow(unreachable_code)]
7837fn dot_q4_row_i8(packed: &[u8], scales: &[u8], g0: usize, gpr: usize, xq: &[i8]) -> f32 {
7838    #[cfg(target_arch = "aarch64")]
7839    unsafe {
7840        return dot_q4_row_sdot(packed, scales, g0, gpr, xq);
7841    }
7842    #[cfg(target_arch = "x86_64")]
7843    unsafe {
7844        return dot_q4_row_avx2(packed, scales, g0, gpr, xq);
7845    }
7846    let mut acc = 0f32;
7847    for gi in 0..gpr {
7848        let g = g0 + gi;
7849        let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
7850        let mut d = 0i32;
7851        for (k, &b) in packed[g * 16..(g + 1) * 16].iter().enumerate() {
7852            d += ((b & 0x0F) as i32 - 8) * xq[gi * GROUP_SIZE + k * 2] as i32
7853                + (((b >> 4) & 0x0F) as i32 - 8) * xq[gi * GROUP_SIZE + k * 2 + 1] as i32;
7854        }
7855        acc += d as f32 * s;
7856    }
7857    acc
7858}
7859
7860/// Two-activation q4 row via the A8W8 int8 path (see `dot_q4_row_i8`).
7861#[inline]
7862#[allow(unreachable_code)]
7863fn dot_q4_row_i8_2(
7864    packed: &[u8],
7865    scales: &[u8],
7866    g0: usize,
7867    gpr: usize,
7868    xq1: &[i8],
7869    xq2: &[i8],
7870) -> (f32, f32) {
7871    #[cfg(target_arch = "aarch64")]
7872    unsafe {
7873        return dot_q4_row_sdot2(packed, scales, g0, gpr, xq1, xq2);
7874    }
7875    #[cfg(target_arch = "x86_64")]
7876    unsafe {
7877        return dot_q4_row_avx2_2(packed, scales, g0, gpr, xq1, xq2);
7878    }
7879    (
7880        dot_q4_row_i8(packed, scales, g0, gpr, xq1),
7881        dot_q4_row_i8(packed, scales, g0, gpr, xq2),
7882    )
7883}
7884
7885/// One q4 row range via SDOT (kernel body of `q4matvec`, extracted so
7886/// multi-matrix jobs can drive it for several tensors in one dispatch).
7887#[allow(clippy::too_many_arguments)]
7888fn q4_range_a8w8(
7889    packed: &[u8],
7890    scales: &[u8],
7891    gpr: usize,
7892    cols: usize,
7893    act: &SplitAct,
7894    out: SendMut,
7895    start: usize,
7896    end: usize,
7897) {
7898    for r in start..end {
7899        let mut acc = dot_q4_row_i8(packed, scales, r * gpr, gpr, &act.xq) * act.sx;
7900        // xq is zeroed at outlier slots — add the exact terms.
7901        for &(j, xv) in &act.outliers {
7902            let flat = r * cols + j;
7903            let byte = packed[flat / 2];
7904            let nib = if flat & 1 == 0 {
7905                byte & 0x0F
7906            } else {
7907                byte >> 4
7908            };
7909            let s = f16_to_f32(u16::from_le_bytes([
7910                scales[(flat / GROUP_SIZE) * 2],
7911                scales[(flat / GROUP_SIZE) * 2 + 1],
7912            ]));
7913            acc += ((nib as i32 - 8) as f32) * s * xv;
7914        }
7915        // SAFETY: disjoint row ranges per worker.
7916        unsafe { *out.at(r) = acc };
7917    }
7918}
7919
7920/// Two-input q4 row range via the A8W8 int8 path — kernel body of
7921/// `q4matvec2`, extracted for pair multi-matrix jobs.
7922#[allow(clippy::too_many_arguments)]
7923fn q4_range2_a8w8(
7924    packed: &[u8],
7925    scales: &[u8],
7926    gpr: usize,
7927    cols: usize,
7928    a1: &SplitAct,
7929    a2: &SplitAct,
7930    p1: SendMut,
7931    p2: SendMut,
7932    start: usize,
7933    end: usize,
7934) {
7935    for r in start..end {
7936        let (s1, s2) = dot_q4_row_i8_2(packed, scales, r * gpr, gpr, &a1.xq, &a2.xq);
7937        let mut acc1 = s1 * a1.sx;
7938        let mut acc2 = s2 * a2.sx;
7939        // xq is zeroed at outlier slots — add the exact terms.
7940        let fix = |outliers: &[(usize, f32)], acc: &mut f32| {
7941            for &(j, xv) in outliers {
7942                let flat = r * cols + j;
7943                let byte = packed[flat / 2];
7944                let nib = if flat & 1 == 0 {
7945                    byte & 0x0F
7946                } else {
7947                    byte >> 4
7948                };
7949                let s = f16_to_f32(u16::from_le_bytes([
7950                    scales[(flat / GROUP_SIZE) * 2],
7951                    scales[(flat / GROUP_SIZE) * 2 + 1],
7952                ]));
7953                *acc += ((nib as i32 - 8) as f32) * s * xv;
7954            }
7955        };
7956        fix(&a1.outliers, &mut acc1);
7957        fix(&a2.outliers, &mut acc2);
7958        // SAFETY: disjoint row ranges per worker.
7959        unsafe {
7960            *p1.at(r) = acc1;
7961            *p2.at(r) = acc2;
7962        }
7963    }
7964}
7965
7966/// Exact scalar q4 row range (same extraction, non-SDOT path).
7967fn q4_range_f32(
7968    packed: &[u8],
7969    scales: &[u8],
7970    gpr: usize,
7971    x: &[f32],
7972    out: SendMut,
7973    start: usize,
7974    end: usize,
7975) {
7976    for r in start..end {
7977        let mut acc = 0f32;
7978        for gi in 0..gpr {
7979            let g = r * gpr + gi;
7980            let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
7981            let pk = &packed[g * 16..(g + 1) * 16];
7982            let xg = &x[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
7983            let mut ga = 0f32;
7984            for (k, &b) in pk.iter().enumerate() {
7985                ga += ((b & 0x0F) as f32 - 8.0) * xg[k * 2]
7986                    + (((b >> 4) & 0x0F) as f32 - 8.0) * xg[k * 2 + 1];
7987            }
7988            acc += ga * s;
7989        }
7990        // SAFETY: disjoint row ranges per worker.
7991        unsafe { *out.at(r) = acc };
7992    }
7993}
7994
7995/// Fused two-input q4 matvec: nibbles are unpacked ONCE per group and
7996/// dotted against both activations (was: two full matvecs — double
7997/// weight traffic). Per-lane math matches `q4matvec` exactly.
7998#[allow(clippy::too_many_arguments)]
7999fn q4matvec2(
8000    bytes: &[u8],
8001    x1: &[f32],
8002    x2: &[f32],
8003    rows: usize,
8004    cols: usize,
8005    o1: &mut [f32],
8006    o2: &mut [f32],
8007    pool: Option<&Pool>,
8008) {
8009    debug_assert_eq!(o1.len(), rows);
8010    debug_assert_eq!(o2.len(), rows);
8011    let (packed, scales) = q4_split(bytes, rows, cols);
8012    let gpr = cols / GROUP_SIZE;
8013
8014    if a8w8_enabled() {
8015        let a1 = split_act(x1);
8016        let a2 = split_act(x2);
8017        let p1 = SendMut(o1.as_mut_ptr());
8018        let p2 = SendMut(o2.as_mut_ptr());
8019        let run = move |start: usize, end: usize| {
8020            q4_range2_a8w8(packed, scales, gpr, cols, &a1, &a2, p1, p2, start, end)
8021        };
8022        dispatch_rows(pool, rows, &run);
8023        return;
8024    }
8025
8026    let p1 = SendMut(o1.as_mut_ptr());
8027    let p2 = SendMut(o2.as_mut_ptr());
8028    let run = move |start: usize, end: usize| {
8029        q4_range2_f32(packed, scales, gpr, x1, x2, p1, p2, start, end)
8030    };
8031    dispatch_rows(pool, rows, &run);
8032}
8033
8034/// Two-input exact scalar q4 row range (same extraction).
8035#[allow(clippy::too_many_arguments)]
8036fn q4_range2_f32(
8037    packed: &[u8],
8038    scales: &[u8],
8039    gpr: usize,
8040    x1: &[f32],
8041    x2: &[f32],
8042    p1: SendMut,
8043    p2: SendMut,
8044    start: usize,
8045    end: usize,
8046) {
8047    for r in start..end {
8048        let (mut acc1, mut acc2) = (0f32, 0f32);
8049        for gi in 0..gpr {
8050            let g = r * gpr + gi;
8051            let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
8052            let pk = &packed[g * 16..(g + 1) * 16];
8053            let x1g = &x1[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
8054            let x2g = &x2[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
8055            let (mut g1, mut g2) = (0f32, 0f32);
8056            for (k, &b) in pk.iter().enumerate() {
8057                let wl = (b & 0x0F) as f32 - 8.0;
8058                let wh = ((b >> 4) & 0x0F) as f32 - 8.0;
8059                g1 += wl * x1g[k * 2] + wh * x1g[k * 2 + 1];
8060                g2 += wl * x2g[k * 2] + wh * x2g[k * 2 + 1];
8061            }
8062            acc1 += g1 * s;
8063            acc2 += g2 * s;
8064        }
8065        // SAFETY: disjoint row ranges per worker.
8066        unsafe {
8067            *p1.at(r) = acc1;
8068            *p2.at(r) = acc2;
8069        }
8070    }
8071}
8072
8073thread_local! {
8074    /// Per-worker decoded-row scratch for the batched q4/vbit kernels
8075    /// (centered i8 for SDOT, f32 for the exact/scalar paths).
8076    static ROW_I8: std::cell::RefCell<Vec<u8>> = const { std::cell::RefCell::new(Vec::new()) };
8077    static ROW_F32: std::cell::RefCell<Vec<f32>> = const { std::cell::RefCell::new(Vec::new()) };
8078}
8079
8080/// Batched q4 matmat: each weight row is unpacked from the mmap ONCE
8081/// and dotted against ALL b activations (prefill used to fall back to b
8082/// full matvecs — b× weight traffic and b× nibble decode). Per-position
8083/// math matches `q4matvec` exactly: same group order, same accumulation.
8084/// `out` is row-major [b, rows] like `qmatmat`.
8085#[allow(clippy::too_many_arguments)]
8086fn q4matmat(
8087    bytes: &[u8],
8088    xs_all: &[f32],
8089    b: usize,
8090    rows: usize,
8091    cols: usize,
8092    out: &mut [f32],
8093    pool: Option<&Pool>,
8094) {
8095    debug_assert_eq!(xs_all.len(), b * cols);
8096    debug_assert_eq!(out.len(), b * rows);
8097    let (packed, scales) = q4_split(bytes, rows, cols);
8098    let gpr = cols / GROUP_SIZE;
8099    let gscale = |g: usize| f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
8100
8101    if a8w8_enabled() {
8102        let acts: Vec<SplitAct> = (0..b)
8103            .map(|bi| split_act(&xs_all[bi * cols..(bi + 1) * cols]))
8104            .collect();
8105        let acts = &acts;
8106        let out_addr = SendMut(out.as_mut_ptr());
8107        let run = move |start: usize, end: usize| {
8108            ROW_I8.with(|rb| {
8109                let mut buf = rb.borrow_mut();
8110                buf.resize(cols, 0);
8111                for r in start..end {
8112                    // Unpack the row's nibbles to centered i8 once
8113                    // (element 2k = low nibble, 2k+1 = high — flat order,
8114                    // same as dot_q4_row_sdot's zip).
8115                    for gi in 0..gpr {
8116                        let g = r * gpr + gi;
8117                        for (k, &bt) in packed[g * 16..(g + 1) * 16].iter().enumerate() {
8118                            buf[gi * GROUP_SIZE + k * 2] = ((bt & 0x0F) as i32 - 8) as i8 as u8;
8119                            buf[gi * GROUP_SIZE + k * 2 + 1] =
8120                                (((bt >> 4) & 0x0F) as i32 - 8) as i8 as u8;
8121                        }
8122                    }
8123                    let mut bi = 0usize;
8124                    #[cfg(target_arch = "x86_64")]
8125                    if avx2_enabled() && blocked_enabled() {
8126                        while bi + 4 <= acts.len() {
8127                            let xs = [
8128                                acts[bi].xq.as_slice(),
8129                                acts[bi + 1].xq.as_slice(),
8130                                acts[bi + 2].xq.as_slice(),
8131                                acts[bi + 3].xq.as_slice(),
8132                            ];
8133                            let d = unsafe {
8134                                if vnni_tiles_enabled() {
8135                                    dot_q4b_row_1x4_vnni(&buf, scales, r * gpr, gpr, xs)
8136                                } else {
8137                                    dot_q4b_row_1x4_avx2(&buf, scales, r * gpr, gpr, xs)
8138                                }
8139                            };
8140                            for k in 0..4 {
8141                                let act = &acts[bi + k];
8142                                let mut acc = d[k] * act.sx;
8143                                for &(j, xv) in &act.outliers {
8144                                    acc += (buf[j] as i8) as f32
8145                                        * gscale((r * cols + j) / GROUP_SIZE)
8146                                        * xv;
8147                                }
8148                                // SAFETY: disjoint (bi, r) cells per worker.
8149                                unsafe { *out_addr.at((bi + k) * rows + r) = acc };
8150                            }
8151                            bi += 4;
8152                        }
8153                    }
8154                    while bi < acts.len() {
8155                        let act = &acts[bi];
8156                        let mut acc = 0f32;
8157                        for gi in 0..gpr {
8158                            let d = dot_i8_i8(
8159                                &buf[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE],
8160                                &act.xq[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE],
8161                            );
8162                            acc += d as f32 * gscale(r * gpr + gi);
8163                        }
8164                        acc *= act.sx;
8165                        // xq is zeroed at outlier slots — exact terms.
8166                        for &(j, xv) in &act.outliers {
8167                            acc += (buf[j] as i8) as f32 * gscale((r * cols + j) / GROUP_SIZE) * xv;
8168                        }
8169                        // SAFETY: disjoint (bi, r) cells per worker row range.
8170                        unsafe { *out_addr.at(bi * rows + r) = acc };
8171                        bi += 1;
8172                    }
8173                }
8174            })
8175        };
8176        dispatch_rows(pool, rows, &run);
8177        return;
8178    }
8179
8180    let out_addr = SendMut(out.as_mut_ptr());
8181    let run = move |start: usize, end: usize| {
8182        ROW_F32.with(|rb| {
8183            let mut buf = rb.borrow_mut();
8184            buf.resize(cols, 0.0);
8185            for r in start..end {
8186                // Decode raw (nib − 8) values once; scales stay per-group
8187                // so the accumulation order matches q4matvec bit-for-bit.
8188                for gi in 0..gpr {
8189                    let g = r * gpr + gi;
8190                    for (k, &bt) in packed[g * 16..(g + 1) * 16].iter().enumerate() {
8191                        buf[gi * GROUP_SIZE + k * 2] = (bt & 0x0F) as f32 - 8.0;
8192                        buf[gi * GROUP_SIZE + k * 2 + 1] = ((bt >> 4) & 0x0F) as f32 - 8.0;
8193                    }
8194                }
8195                for bi in 0..b {
8196                    let x = &xs_all[bi * cols..(bi + 1) * cols];
8197                    let mut acc = 0f32;
8198                    for gi in 0..gpr {
8199                        let mut ga = 0f32;
8200                        // Pairwise (lo + hi) addition, matching
8201                        // q4matvec's `ga += lo·x + hi·x` shape exactly —
8202                        // a flat one-per-element loop rounds differently
8203                        // and broke bit-parity on the scalar (x86) path.
8204                        for k in 0..GROUP_SIZE / 2 {
8205                            let e = gi * GROUP_SIZE + k * 2;
8206                            ga += buf[e] * x[e] + buf[e + 1] * x[e + 1];
8207                        }
8208                        acc += ga * gscale(r * gpr + gi);
8209                    }
8210                    // SAFETY: disjoint (bi, r) cells per worker row range.
8211                    unsafe { *out_addr.at(bi * rows + r) = acc };
8212                }
8213            }
8214        })
8215    };
8216    dispatch_rows(pool, rows, &run);
8217}
8218
8219/// Batched vbit matmat: each variable-bit row is decoded from the mmap
8220/// ONCE for the whole microbatch. Same per-position math as
8221/// `vbitmatvec` (SDOT A8W8 with exact outliers / exact f32 for b=8 rows
8222/// and the scalar path).
8223#[allow(clippy::too_many_arguments)]
8224fn vbitmatmat(
8225    bytes: &[u8],
8226    offsets: &[usize],
8227    xs_all: &[f32],
8228    b: usize,
8229    rows: usize,
8230    cols: usize,
8231    out: &mut [f32],
8232    pool: Option<&Pool>,
8233) {
8234    debug_assert_eq!(xs_all.len(), b * cols);
8235    debug_assert_eq!(out.len(), b * rows);
8236    debug_assert_eq!(offsets.len(), rows + 1);
8237    let ng = cols / GROUP_SIZE;
8238    let bits = &bytes[..rows];
8239    let sc_off = rows;
8240    let gscale = |r: usize, g: usize| {
8241        let so = (r * ng + g) * 2;
8242        f16_to_f32(u16::from_le_bytes([
8243            bytes[sc_off + so],
8244            bytes[sc_off + so + 1],
8245        ]))
8246    };
8247
8248    // Decode row r's raw (u − L) values into `dst` (f32, unscaled).
8249    let decode_f32 = |r: usize, dst: &mut [f32]| {
8250        let bw = bits[r] as usize;
8251        let l = ((1i32 << (bw - 1)) - 1) as f32;
8252        let data = &bytes[offsets[r]..offsets[r + 1]];
8253        let (mut acc, mut nbits, mut idx) = (0u64, 0usize, 0usize);
8254        for d in dst.iter_mut() {
8255            while nbits < bw {
8256                acc = (acc << 8) | data[idx] as u64;
8257                idx += 1;
8258                nbits += 8;
8259            }
8260            let u = ((acc >> (nbits - bw)) & ((1u64 << bw) - 1)) as f32;
8261            nbits -= bw;
8262            *d = u - l;
8263        }
8264    };
8265
8266    if a8w8_enabled() {
8267        let acts: Vec<SplitAct> = (0..b)
8268            .map(|bi| split_act(&xs_all[bi * cols..(bi + 1) * cols]))
8269            .collect();
8270        let acts = &acts;
8271        let out_addr = SendMut(out.as_mut_ptr());
8272        let run = move |start: usize, end: usize| {
8273            for r in start..end {
8274                let bw = bits[r] as usize;
8275                if bw == 8 {
8276                    // u−L reaches 128 → no i8 path; decode once, exact
8277                    // f32 dots for every position (same as vbitmatvec).
8278                    ROW_F32.with(|rb| {
8279                        let mut buf = rb.borrow_mut();
8280                        buf.resize(cols, 0.0);
8281                        decode_f32(r, &mut buf);
8282                        for bi in 0..b {
8283                            let x = &xs_all[bi * cols..(bi + 1) * cols];
8284                            let mut dot = 0f32;
8285                            for g in 0..ng {
8286                                let mut gd = 0f32;
8287                                for k in 0..GROUP_SIZE {
8288                                    gd += buf[g * GROUP_SIZE + k] * x[g * GROUP_SIZE + k];
8289                                }
8290                                dot += gd * gscale(r, g);
8291                            }
8292                            // SAFETY: disjoint (bi, r) cells per worker range.
8293                            unsafe { *out_addr.at(bi * rows + r) = dot };
8294                        }
8295                    });
8296                    continue;
8297                }
8298                let l = (1i32 << (bw - 1)) - 1;
8299                let data = &bytes[offsets[r]..offsets[r + 1]];
8300                ROW_I8.with(|rb| {
8301                    let mut buf = rb.borrow_mut();
8302                    buf.resize(cols, 0);
8303                    #[inline(always)]
8304                    fn fill<const B: usize>(data: &[u8], l: i32, buf: &mut [u8]) {
8305                        for (blk, chunk) in buf.chunks_exact_mut(8).enumerate() {
8306                            let u = unpack8::<B>(&data[blk * B..]);
8307                            for k in 0..8 {
8308                                chunk[k] = (u[k] - l) as i8 as u8;
8309                            }
8310                        }
8311                    }
8312                    match bw {
8313                        3 => fill::<3>(data, l, &mut buf),
8314                        4 => vbit_fill4(data, &mut buf),
8315                        5 => fill::<5>(data, l, &mut buf),
8316                        6 => fill::<6>(data, l, &mut buf),
8317                        _ => unreachable!("vbit bit-width {bw} (validated at load)"),
8318                    }
8319                    let mut bi = 0usize;
8320                    // The vbit scale table shares q4_block's layout
8321                    // (contiguous f16 per (row·ng + g)), so the same
8322                    // blocked 1×4 kernel serves the decoded row.
8323                    #[cfg(target_arch = "x86_64")]
8324                    if avx2_enabled() && blocked_enabled() {
8325                        while bi + 4 <= acts.len() {
8326                            let xs = [
8327                                acts[bi].xq.as_slice(),
8328                                acts[bi + 1].xq.as_slice(),
8329                                acts[bi + 2].xq.as_slice(),
8330                                acts[bi + 3].xq.as_slice(),
8331                            ];
8332                            let sxs = [
8333                                acts[bi].sx,
8334                                acts[bi + 1].sx,
8335                                acts[bi + 2].sx,
8336                                acts[bi + 3].sx,
8337                            ];
8338                            let d = unsafe {
8339                                if vnni_tiles_enabled() {
8340                                    dot_q4b_row_1x4_sx_vnni(
8341                                        &buf,
8342                                        &bytes[sc_off..],
8343                                        r * ng,
8344                                        ng,
8345                                        xs,
8346                                        sxs,
8347                                    )
8348                                } else {
8349                                    dot_q4b_row_1x4_sx_avx2(
8350                                        &buf,
8351                                        &bytes[sc_off..],
8352                                        r * ng,
8353                                        ng,
8354                                        xs,
8355                                        sxs,
8356                                    )
8357                                }
8358                            };
8359                            for k in 0..4 {
8360                                let act = &acts[bi + k];
8361                                let mut dot = d[k];
8362                                for &(j, xv) in &act.outliers {
8363                                    dot += (buf[j] as i8) as f32 * gscale(r, j / GROUP_SIZE) * xv;
8364                                }
8365                                // SAFETY: disjoint (bi, r) cells per worker.
8366                                unsafe { *out_addr.at((bi + k) * rows + r) = dot };
8367                            }
8368                            bi += 4;
8369                        }
8370                    }
8371                    while bi < acts.len() {
8372                        let act = &acts[bi];
8373                        let mut dot = 0f32;
8374                        for g in 0..ng {
8375                            let d = dot_i8_i8(
8376                                &buf[g * GROUP_SIZE..(g + 1) * GROUP_SIZE],
8377                                &act.xq[g * GROUP_SIZE..(g + 1) * GROUP_SIZE],
8378                            ) as f32
8379                                * act.sx;
8380                            dot += d * gscale(r, g);
8381                        }
8382                        for &(j, xv) in &act.outliers {
8383                            dot += (buf[j] as i8) as f32 * gscale(r, j / GROUP_SIZE) * xv;
8384                        }
8385                        // SAFETY: disjoint (bi, r) cells per worker range.
8386                        unsafe { *out_addr.at(bi * rows + r) = dot };
8387                        bi += 1;
8388                    }
8389                });
8390            }
8391        };
8392        dispatch_rows(pool, rows, &run);
8393        return;
8394    }
8395
8396    let out_addr = SendMut(out.as_mut_ptr());
8397    let run = move |start: usize, end: usize| {
8398        ROW_F32.with(|rb| {
8399            let mut buf = rb.borrow_mut();
8400            buf.resize(cols, 0.0);
8401            for r in start..end {
8402                decode_f32(r, &mut buf);
8403                for bi in 0..b {
8404                    let x = &xs_all[bi * cols..(bi + 1) * cols];
8405                    let mut dot = 0f32;
8406                    for g in 0..ng {
8407                        let mut gd = 0f32;
8408                        for k in 0..GROUP_SIZE {
8409                            gd += buf[g * GROUP_SIZE + k] * x[g * GROUP_SIZE + k];
8410                        }
8411                        dot += gd * gscale(r, g);
8412                    }
8413                    // SAFETY: disjoint (bi, r) cells per worker range.
8414                    unsafe { *out_addr.at(bi * rows + r) = dot };
8415                }
8416            }
8417        })
8418    };
8419    dispatch_rows(pool, rows, &run);
8420}
8421
8422/// Build a GPU batch job for a q8-family mapped tensor (primary
8423/// shard): prescaled input + directory coordinates. None → not
8424/// GPU-eligible, caller stays on the CPU.
8425pub(crate) fn gpu_batch_job<'a>(
8426    t: &'a QTensor,
8427    x: &[f32],
8428) -> Option<(std::sync::Arc<CmfModel>, crate::gpu::BatchJob<'a>)> {
8429    match t {
8430        QTensor::Mapped {
8431            model,
8432            idx,
8433            dtype: dt @ (TensorDtype::Q8Row | TensorDtype::Q8_2f),
8434            rows,
8435            cols,
8436            row_scale,
8437            col_field,
8438            ..
8439        } => Some((
8440            model.clone(),
8441            crate::gpu::BatchJob {
8442                idx: *idx,
8443                rows: *rows,
8444                cols: *cols,
8445                row_scale,
8446                xs: prescale(x, col_field, *dt).into_owned(),
8447                layout: crate::gpu::BatchLayout::Q8,
8448            },
8449        )),
8450        // q1: raw f32 activations, tile-embedded scales.
8451        QTensor::Mapped {
8452            model,
8453            idx,
8454            dtype: TensorDtype::Q1,
8455            rows,
8456            cols,
8457            ..
8458        } => Some((
8459            model.clone(),
8460            crate::gpu::BatchJob {
8461                idx: *idx,
8462                rows: *rows,
8463                cols: *cols,
8464                row_scale: &[],
8465                xs: x.to_vec(),
8466                layout: crate::gpu::BatchLayout::Q1,
8467            },
8468        )),
8469        // q4_tiled / q4tp: raw f32 activations; the scales live in the
8470        // payload (inline tiles / row ladder), so row_scale stays empty.
8471        // The GDN projection batch already runs these layouts on Metal —
8472        // this arm lets the attention QKV batch reach the same kernels.
8473        QTensor::Mapped {
8474            model,
8475            idx,
8476            dtype: dt @ (TensorDtype::Q4Tiled | TensorDtype::Q4TiledP),
8477            rows,
8478            cols,
8479            ..
8480        } => Some((
8481            model.clone(),
8482            crate::gpu::BatchJob {
8483                idx: *idx,
8484                rows: *rows,
8485                cols: *cols,
8486                row_scale: &[],
8487                xs: x.to_vec(),
8488                layout: if *dt == TensorDtype::Q4Tiled {
8489                    crate::gpu::BatchLayout::Q4t
8490                } else {
8491                    crate::gpu::BatchLayout::Q4tp
8492                },
8493            },
8494        )),
8495        _ => None,
8496    }
8497}
8498
8499thread_local! {
8500    static PRESCALE_BUF1: std::cell::RefCell<Vec<f32>> = const { std::cell::RefCell::new(Vec::new()) };
8501    static PRESCALE_BUF2: std::cell::RefCell<Vec<f32>> = const { std::cell::RefCell::new(Vec::new()) };
8502}
8503
8504pub(crate) fn prescale<'a>(
8505    x: &'a [f32],
8506    col_field: &[f32],
8507    dtype: TensorDtype,
8508) -> std::borrow::Cow<'a, [f32]> {
8509    if dtype == TensorDtype::Q8_2f {
8510        x.iter().zip(col_field).map(|(a, c)| a * c).collect()
8511    } else {
8512        std::borrow::Cow::Borrowed(x)
8513    }
8514}
8515
8516/// θ col-field fold for q8_2f activations. Borrowed pass-through for
8517/// every other dtype, using thread-local buffers to eliminate per-matvec allocations.
8518pub(crate) fn prescale_with<R, F: FnOnce(&[f32]) -> R>(
8519    x: &[f32],
8520    col_field: &[f32],
8521    dtype: TensorDtype,
8522    buf_id: u8,
8523    f: F,
8524) -> R {
8525    if dtype == TensorDtype::Q8_2f {
8526        if buf_id == 1 {
8527            PRESCALE_BUF1.with(|b| {
8528                let mut buf = b.borrow_mut();
8529                buf.clear();
8530                buf.extend(x.iter().zip(col_field).map(|(a, c)| a * c));
8531                f(&buf)
8532            })
8533        } else {
8534            PRESCALE_BUF2.with(|b| {
8535                let mut buf = b.borrow_mut();
8536                buf.clear();
8537                buf.extend(x.iter().zip(col_field).map(|(a, c)| a * c));
8538                f(&buf)
8539            })
8540        }
8541    } else {
8542        f(x)
8543    }
8544}
8545
8546// ───────────────────── x86-64 AVX2 kernels (roadmap этап 2) ─────────────────────
8547
8548/// AVX2+FMA available? Default ON when the CPU supports both;
8549/// `CMF_AVX2=0` disables (falls back to the autovectorized loops).
8550#[cfg(target_arch = "x86_64")]
8551pub(crate) fn avx2_enabled() -> bool {
8552    use std::sync::OnceLock;
8553    static ON: OnceLock<bool> = OnceLock::new();
8554    *ON.get_or_init(|| {
8555        std::env::var("CMF_AVX2").map(|v| v != "0").unwrap_or(true)
8556            && std::arch::is_x86_feature_detected!("avx2")
8557            && std::arch::is_x86_feature_detected!("fma")
8558    })
8559}
8560
8561/// AVX2 A8W8 allowed? The quantized-activation contract is switched by
8562/// the SAME env as the ARM SDOT path: `CMF_SDOT=0` keeps exact kernels
8563/// (the golden-parity exact gate relies on it) — AVX2 f32 kernels stay
8564/// active either way, they are exact (regrouped sums only).
8565#[cfg(target_arch = "x86_64")]
8566fn avx2_a8w8_enabled() -> bool {
8567    use std::sync::OnceLock;
8568    static ON: OnceLock<bool> = OnceLock::new();
8569    *ON.get_or_init(|| {
8570        avx2_enabled() && std::env::var("CMF_SDOT").map(|v| v != "0").unwrap_or(true)
8571    })
8572}
8573
8574/// A8W8 quantized-activation path available on THIS machine? One
8575/// switch across architectures: ARM dotprod (CMF_SDOT) or x86 AVX2
8576/// (CMF_AVX2 + the same CMF_SDOT exact-contract override).
8577#[inline]
8578pub(crate) fn a8w8_enabled() -> bool {
8579    #[cfg(target_arch = "aarch64")]
8580    {
8581        sdot_enabled()
8582    }
8583    #[cfg(target_arch = "x86_64")]
8584    {
8585        avx2_a8w8_enabled()
8586    }
8587    #[cfg(not(any(target_arch = "aarch64", target_arch = "x86_64")))]
8588    {
8589        false
8590    }
8591}
8592
8593/// int8·int8 dot dispatch: SDOT on ARM; AVX-512 VNNI (vpdpbusd) or AVX2
8594/// maddubs on x86. Callers are gated by `a8w8_enabled()`.
8595#[inline]
8596#[allow(unreachable_code)]
8597fn dot_i8_i8(w: &[u8], xq: &[i8]) -> i32 {
8598    #[cfg(target_arch = "aarch64")]
8599    unsafe {
8600        return dot_i8_sdot(w, xq);
8601    }
8602    #[cfg(target_arch = "x86_64")]
8603    unsafe {
8604        if avx512vnni_enabled() {
8605            return dot_i8_i8_vnni(w, xq);
8606        }
8607        return dot_i8_i8_avx2(w, xq);
8608    }
8609    w.iter()
8610        .zip(xq)
8611        .map(|(&a, &b)| (a as i8) as i32 * b as i32)
8612        .sum()
8613}
8614
8615/// AVX-512 VNNI available? (F+BW+VL+VNNI; `CMF_AVX512=0` falls back to
8616/// AVX2.) VL matters: short 32-byte groups (q4/vbit) ride the 256-bit
8617/// `vpdpbusd` encoding.
8618#[cfg(target_arch = "x86_64")]
8619fn avx512vnni_enabled() -> bool {
8620    use std::sync::OnceLock;
8621    static ON: OnceLock<bool> = OnceLock::new();
8622    *ON.get_or_init(|| {
8623        std::env::var("CMF_AVX512")
8624            .map(|v| v != "0")
8625            .unwrap_or(true)
8626            && std::arch::is_x86_feature_detected!("avx512f")
8627            && std::arch::is_x86_feature_detected!("avx512bw")
8628            && std::arch::is_x86_feature_detected!("avx512vl")
8629            && std::arch::is_x86_feature_detected!("avx512vnni")
8630    })
8631}
8632
8633/// Grouped-codec VNNI arms (the q4t/q4b/q1/q1t tile kernels): default
8634/// ON where AVX-512 VNNI exists (`CMF_VNNI_TILES=0` opt-out). Measured
8635/// on Ryzen 7950X (Zen4, 3 alternating process pairs, blocked GEMM
8636/// 4864×896 b=256): q4t 63→68 GF/s (+8%), q1 53→56 (+6%), q4b 72→75
8637/// (+4%) — consistent, no leg regressed. The tile kernels keep a
8638/// horizontal reduce per 32-weight group, so the `vpdpbusd` saving is
8639/// smaller than the long-dot q8 win (+13%), but it is real and free.
8640#[cfg(target_arch = "x86_64")]
8641fn vnni_tiles_enabled() -> bool {
8642    use std::sync::OnceLock;
8643    static ON: OnceLock<bool> = OnceLock::new();
8644    *ON.get_or_init(|| {
8645        std::env::var("CMF_VNNI_TILES")
8646            .map(|v| v != "0")
8647            .unwrap_or(true)
8648            && avx512vnni_enabled()
8649    })
8650}
8651
8652/// One 256-bit u8×i8 dot → i32 via `vpdpbusd` into a fresh accumulator
8653/// plus the same horizontal reduce the AVX2 kernels use. Products are
8654/// bounded (|w| ≤ 8 or ≤ 1), so maddubs never saturated — the i32 sum
8655/// is bit-identical to the maddubs+madd pair it replaces.
8656#[cfg(target_arch = "x86_64")]
8657#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
8658#[inline]
8659unsafe fn dpbusd_hsum(aw: core::arch::x86_64::__m256i, xs: core::arch::x86_64::__m256i) -> i32 {
8660    // SAFETY: pure register math.
8661    unsafe {
8662        use core::arch::x86_64::*;
8663        let d = _mm256_dpbusd_epi32(_mm256_setzero_si256(), aw, xs);
8664        let hi128 = _mm256_extracti128_si256::<1>(d);
8665        let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
8666        let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
8667        let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
8668        _mm_cvtsi128_si32(s32)
8669    }
8670}
8671
8672/// int8·int8 via AVX-512 VNNI: `vpdpbusd` fuses the maddubs+madd+add
8673/// triple into one u8×i8 dot-accumulate. AVX-512 has no vpsignb, so the
8674/// |w|·sign(x,w) trick becomes |w| × (x negated where w<0) via a mask
8675/// subtract — w==0 lanes contribute 0 through |w|=0 either way.
8676#[cfg(target_arch = "x86_64")]
8677#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
8678unsafe fn dot_i8_i8_vnni(w: &[u8], xq: &[i8]) -> i32 {
8679    // SAFETY: callers uphold slice-length contracts (see call sites).
8680    unsafe {
8681        use core::arch::x86_64::*;
8682        let n = w.len();
8683        let mut j = 0usize;
8684        let mut total: i32;
8685        // 4 independent accumulators: vpdpbusd is its own loop-carried
8686        // dependency (~5-cycle latency) — a single-acc loop runs
8687        // latency-bound and LOSES to the AVX2 maddubs kernel, measured
8688        // on Granite Rapids.
8689        {
8690            #[inline(always)]
8691            unsafe fn step(
8692                w: *const u8,
8693                x: *const i8,
8694                acc: core::arch::x86_64::__m512i,
8695            ) -> core::arch::x86_64::__m512i {
8696                unsafe {
8697                    use core::arch::x86_64::*;
8698                    let wv = _mm512_loadu_si512(w as *const _);
8699                    let xv = _mm512_loadu_si512(x as *const _);
8700                    let aw = _mm512_abs_epi8(wv);
8701                    let neg = _mm512_movepi8_mask(wv);
8702                    let sx = _mm512_mask_sub_epi8(xv, neg, _mm512_setzero_si512(), xv);
8703                    _mm512_dpbusd_epi32(acc, aw, sx)
8704                }
8705            }
8706            let (mut a0, mut a1, mut a2, mut a3) = (
8707                _mm512_setzero_si512(),
8708                _mm512_setzero_si512(),
8709                _mm512_setzero_si512(),
8710                _mm512_setzero_si512(),
8711            );
8712            while j + 256 <= n {
8713                a0 = step(w.as_ptr().add(j), xq.as_ptr().add(j), a0);
8714                a1 = step(w.as_ptr().add(j + 64), xq.as_ptr().add(j + 64), a1);
8715                a2 = step(w.as_ptr().add(j + 128), xq.as_ptr().add(j + 128), a2);
8716                a3 = step(w.as_ptr().add(j + 192), xq.as_ptr().add(j + 192), a3);
8717                j += 256;
8718            }
8719            while j + 64 <= n {
8720                a0 = step(w.as_ptr().add(j), xq.as_ptr().add(j), a0);
8721                j += 64;
8722            }
8723            let s01 = _mm512_add_epi32(a0, a1);
8724            let s23 = _mm512_add_epi32(a2, a3);
8725            total = _mm512_reduce_add_epi32(_mm512_add_epi32(s01, s23));
8726        }
8727        // 32-wide (q4/vbit groups are exactly 32 bytes).
8728        if j + 32 <= n {
8729            let wv = _mm256_loadu_si256(w.as_ptr().add(j) as *const __m256i);
8730            let xv = _mm256_loadu_si256(xq.as_ptr().add(j) as *const __m256i);
8731            let d = _mm256_dpbusd_epi32(
8732                _mm256_setzero_si256(),
8733                _mm256_abs_epi8(wv),
8734                _mm256_sign_epi8(xv, wv),
8735            );
8736            let hi128 = _mm256_extracti128_si256::<1>(d);
8737            let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
8738            let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
8739            let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
8740            total += _mm_cvtsi128_si32(s32);
8741            j += 32;
8742        }
8743        while j < n {
8744            total += (w[j] as i8) as i32 * xq[j] as i32;
8745            j += 1;
8746        }
8747        total
8748    }
8749}
8750
8751/// i8 row · f32 x via AVX2/FMA (x86 mirror of `dot_i8_f32_neon`).
8752#[cfg(target_arch = "x86_64")]
8753#[target_feature(enable = "avx2,fma")]
8754unsafe fn dot_i8_f32_avx2(w: &[u8], x: &[f32]) -> f32 {
8755    // SAFETY: callers uphold slice-length contracts (see call sites).
8756    unsafe {
8757        use core::arch::x86_64::*;
8758        let n = x.len();
8759        let wp = w.as_ptr();
8760        let xp = x.as_ptr();
8761        let (mut a0, mut a1) = (_mm256_setzero_ps(), _mm256_setzero_ps());
8762        let mut j = 0usize;
8763        while j + 16 <= n {
8764            let wb = _mm_loadu_si128(wp.add(j) as *const __m128i);
8765            let lo = _mm256_cvtepi8_epi32(wb);
8766            let hi = _mm256_cvtepi8_epi32(_mm_srli_si128::<8>(wb));
8767            a0 = _mm256_fmadd_ps(_mm256_cvtepi32_ps(lo), _mm256_loadu_ps(xp.add(j)), a0);
8768            a1 = _mm256_fmadd_ps(_mm256_cvtepi32_ps(hi), _mm256_loadu_ps(xp.add(j + 8)), a1);
8769            j += 16;
8770        }
8771        let acc = _mm256_add_ps(a0, a1);
8772        let hi128 = _mm256_extractf128_ps::<1>(acc);
8773        let s128 = _mm_add_ps(_mm256_castps256_ps128(acc), hi128);
8774        let s64 = _mm_add_ps(s128, _mm_movehl_ps(s128, s128));
8775        let s32 = _mm_add_ss(s64, _mm_shuffle_ps::<1>(s64, s64));
8776        let mut sum = _mm_cvtss_f32(s32);
8777        while j < n {
8778            sum += (*wp.add(j) as i8) as f32 * *xp.add(j);
8779            j += 1;
8780        }
8781        sum
8782    }
8783}
8784
8785/// int8(weight)·int8(activation) → i32 via AVX2 maddubs — the x86
8786/// analogue of the SDOT A8W8 path. `maddubs` takes u8×i8, so the
8787/// standard sign trick applies: |w| × sign(x, w) ≡ w × x per lane.
8788/// Pair saturation is safe: |w|≤128, |x|≤127 → 2·128·127 < 32767.
8789#[cfg(target_arch = "x86_64")]
8790#[target_feature(enable = "avx2")]
8791unsafe fn dot_i8_i8_avx2(w: &[u8], xq: &[i8]) -> i32 {
8792    // SAFETY: callers uphold slice-length contracts (see call sites).
8793    unsafe {
8794        use core::arch::x86_64::*;
8795        let n = w.len();
8796        let ones = _mm256_set1_epi16(1);
8797        let mut acc = _mm256_setzero_si256();
8798        let mut j = 0usize;
8799        while j + 32 <= n {
8800            let wv = _mm256_loadu_si256(w.as_ptr().add(j) as *const __m256i);
8801            let xv = _mm256_loadu_si256(xq.as_ptr().add(j) as *const __m256i);
8802            let p16 = _mm256_maddubs_epi16(_mm256_abs_epi8(wv), _mm256_sign_epi8(xv, wv));
8803            acc = _mm256_add_epi32(acc, _mm256_madd_epi16(p16, ones));
8804            j += 32;
8805        }
8806        let hi128 = _mm256_extracti128_si256::<1>(acc);
8807        let s128 = _mm_add_epi32(_mm256_castsi256_si128(acc), hi128);
8808        let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
8809        let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
8810        let mut s = _mm_cvtsi128_si32(s32);
8811        while j < n {
8812            s += (w[j] as i8) as i32 * xq[j] as i32;
8813            j += 1;
8814        }
8815        s
8816    }
8817}
8818
8819/// smmla 2×4: one instruction covers a 2-row × 2-activation × 8-deep
8820/// tile (32 MACs vs sdot's 16) — the weight pair loads once per 8-k
8821/// slice as a combined 2×8 register and meets two activation pairs.
8822#[cfg(target_arch = "aarch64")]
8823#[target_feature(enable = "neon,i8mm")]
8824unsafe fn dot_i8_smmla_2x4(w0: &[u8], w1: &[u8], xs: [&[i8]; 4]) -> [[i32; 4]; 2] {
8825    // SAFETY: callers uphold slice-length contracts.
8826    unsafe {
8827        use core::arch::aarch64::*;
8828        use core::arch::asm;
8829        let n = w0.len();
8830        let w0p = w0.as_ptr() as *const i8;
8831        let w1p = w1.as_ptr() as *const i8;
8832        // acc01 holds [c(r0,x0) c(r0,x1) c(r1,x0) c(r1,x1)]; acc23 the
8833        // same for x2/x3.
8834        let mut acc01 = vdupq_n_s32(0);
8835        let mut acc23 = vdupq_n_s32(0);
8836        let mut i = 0usize;
8837        while i + 8 <= n {
8838            let wa = vcombine_s8(vld1_s8(w0p.add(i)), vld1_s8(w1p.add(i)));
8839            let xb01 = vcombine_s8(
8840                vld1_s8(xs[0].as_ptr().add(i)),
8841                vld1_s8(xs[1].as_ptr().add(i)),
8842            );
8843            let xb23 = vcombine_s8(
8844                vld1_s8(xs[2].as_ptr().add(i)),
8845                vld1_s8(xs[3].as_ptr().add(i)),
8846            );
8847            asm!(
8848                "smmla {a01:v}.4s, {w:v}.16b, {x01:v}.16b",
8849                "smmla {a23:v}.4s, {w:v}.16b, {x23:v}.16b",
8850                a01 = inout(vreg) acc01, a23 = inout(vreg) acc23,
8851                w = in(vreg) wa, x01 = in(vreg) xb01, x23 = in(vreg) xb23,
8852                options(pure, nomem, nostack),
8853            );
8854            i += 8;
8855        }
8856        let mut out = [[0i32; 4]; 2];
8857        let a01: [i32; 4] = core::mem::transmute(acc01);
8858        let a23: [i32; 4] = core::mem::transmute(acc23);
8859        out[0][0] = a01[0];
8860        out[0][1] = a01[1];
8861        out[1][0] = a01[2];
8862        out[1][1] = a01[3];
8863        out[0][2] = a23[0];
8864        out[0][3] = a23[1];
8865        out[1][2] = a23[2];
8866        out[1][3] = a23[3];
8867        if i < n {
8868            for (k, x) in xs.iter().enumerate() {
8869                for j in i..n {
8870                    out[0][k] += (w0[j] as i8) as i32 * x[j] as i32;
8871                    out[1][k] += (w1[j] as i8) as i32 * x[j] as i32;
8872                }
8873            }
8874        }
8875        out
8876    }
8877}
8878
8879/// ARM twin of the x86 blocked prefill GEMM: two weight rows stay in
8880/// registers across four activation streams, eight sdot accumulators.
8881/// (The per-row form re-read each W row once per activation.)
8882#[cfg(target_arch = "aarch64")]
8883#[target_feature(enable = "neon,dotprod")]
8884unsafe fn dot_i8_sdot_2x4(w0: &[u8], w1: &[u8], xs: [&[i8]; 4]) -> [[i32; 4]; 2] {
8885    // SAFETY: callers uphold slice-length contracts.
8886    unsafe {
8887        use core::arch::aarch64::*;
8888        use core::arch::asm;
8889        let n = w0.len();
8890        let w0p = w0.as_ptr() as *const i8;
8891        let w1p = w1.as_ptr() as *const i8;
8892        let mut acc = [[vdupq_n_s32(0); 4]; 2];
8893        let mut i = 0usize;
8894        while i + 16 <= n {
8895            let wv0 = vld1q_s8(w0p.add(i));
8896            let wv1 = vld1q_s8(w1p.add(i));
8897            for (k, x) in xs.iter().enumerate() {
8898                let xv = vld1q_s8(x.as_ptr().add(i));
8899                let (mut a0, mut a1) = (acc[0][k], acc[1][k]);
8900                asm!(
8901                    "sdot {a0:v}.4s, {w0:v}.16b, {x:v}.16b",
8902                    "sdot {a1:v}.4s, {w1:v}.16b, {x:v}.16b",
8903                    a0 = inout(vreg) a0, a1 = inout(vreg) a1,
8904                    w0 = in(vreg) wv0, w1 = in(vreg) wv1, x = in(vreg) xv,
8905                    options(pure, nomem, nostack),
8906                );
8907                acc[0][k] = a0;
8908                acc[1][k] = a1;
8909            }
8910            i += 16;
8911        }
8912        let mut out = [[0i32; 4]; 2];
8913        for r in 0..2 {
8914            for k in 0..4 {
8915                out[r][k] = vaddvq_s32(acc[r][k]);
8916            }
8917        }
8918        if i < n {
8919            for (k, x) in xs.iter().enumerate() {
8920                for j in i..n {
8921                    out[0][k] += (w0[j] as i8) as i32 * x[j] as i32;
8922                    out[1][k] += (w1[j] as i8) as i32 * x[j] as i32;
8923                }
8924            }
8925        }
8926        out
8927    }
8928}
8929
8930/// Blocked 2 weight rows × 4 activations for the prefill GEMM
8931/// (roadmap P0: packed panels + multi-row accumulators). The two rows'
8932/// abs() live in registers across all four activation streams; the
8933/// sign-fixup is recomputed per pair (the price of the maddubs trick).
8934/// Returns raw i8·i8 dots; the caller applies scales and outliers.
8935#[cfg(target_arch = "x86_64")]
8936#[target_feature(enable = "avx2")]
8937unsafe fn dot_i8_i8_avx2_2x4(w0: &[u8], w1: &[u8], xs: [&[i8]; 4]) -> [[i32; 4]; 2] {
8938    // SAFETY: callers uphold slice-length contracts.
8939    unsafe {
8940        use core::arch::x86_64::*;
8941        let n = w0.len();
8942        let ones = _mm256_set1_epi16(1);
8943        let mut acc = [[_mm256_setzero_si256(); 4]; 2];
8944        let mut j = 0usize;
8945        while j + 32 <= n {
8946            let wv0 = _mm256_loadu_si256(w0.as_ptr().add(j) as *const __m256i);
8947            let wv1 = _mm256_loadu_si256(w1.as_ptr().add(j) as *const __m256i);
8948            let aw0 = _mm256_abs_epi8(wv0);
8949            let aw1 = _mm256_abs_epi8(wv1);
8950            for (k, x) in xs.iter().enumerate() {
8951                let xv = _mm256_loadu_si256(x.as_ptr().add(j) as *const __m256i);
8952                let p0 = _mm256_maddubs_epi16(aw0, _mm256_sign_epi8(xv, wv0));
8953                acc[0][k] = _mm256_add_epi32(acc[0][k], _mm256_madd_epi16(p0, ones));
8954                let p1 = _mm256_maddubs_epi16(aw1, _mm256_sign_epi8(xv, wv1));
8955                acc[1][k] = _mm256_add_epi32(acc[1][k], _mm256_madd_epi16(p1, ones));
8956            }
8957            j += 32;
8958        }
8959        let mut out = [[0i32; 4]; 2];
8960        for r in 0..2 {
8961            for k in 0..4 {
8962                let a = acc[r][k];
8963                let hi128 = _mm256_extracti128_si256::<1>(a);
8964                let s128 = _mm_add_epi32(_mm256_castsi256_si128(a), hi128);
8965                let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
8966                let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
8967                out[r][k] = _mm_cvtsi128_si32(s32);
8968            }
8969        }
8970        if j < n {
8971            for (k, x) in xs.iter().enumerate() {
8972                for i in j..n {
8973                    out[0][k] += (w0[i] as i8) as i32 * x[i] as i32;
8974                    out[1][k] += (w1[i] as i8) as i32 * x[i] as i32;
8975                }
8976            }
8977        }
8978        out
8979    }
8980}
8981
8982/// AVX2/VNNI q8 row dot with exact outlier correction (x86 mirror of
8983/// `row_dot_sdot` — same A8W8 contract). With AVX-512 VNNI the row goes
8984/// through the bias trick: Σ(w+128)·x via pure `vpdpbusd` (no per-lane
8985/// sign fixups), corrected by −128·Σx with Σx precomputed per split.
8986#[cfg(target_arch = "x86_64")]
8987#[inline]
8988fn row_dot_avx2(row: &[u8], act: &SplitAct) -> f32 {
8989    let dot = if avx512vnni_enabled() && row.len() >= 64 {
8990        (unsafe { dot_u8p128_i8_vnni(row, &act.xq) }) - 128 * act.xsum
8991    } else {
8992        unsafe { dot_i8_i8_avx2(row, &act.xq) }
8993    };
8994    let mut acc = dot as f32 * act.sx;
8995    for &(j, xv) in &act.outliers {
8996        acc += (row[j] as i8) as f32 * xv;
8997    }
8998    acc
8999}
9000
9001/// Σ (w[i]+128)·x[i] via pure `vpdpbusd` — the caller subtracts
9002/// 128·Σx. Four independent accumulators (dpbusd is ~5-cycle latency;
9003/// a single-acc loop runs latency-bound, measured on Granite Rapids).
9004#[cfg(target_arch = "x86_64")]
9005#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
9006unsafe fn dot_u8p128_i8_vnni(w: &[u8], xq: &[i8]) -> i32 {
9007    // SAFETY: callers uphold slice-length contracts (see call sites).
9008    unsafe {
9009        use core::arch::x86_64::*;
9010        let n = w.len();
9011        let flip = _mm512_set1_epi8(-128); // XOR 0x80: i8 w → u8 (w+128)
9012        #[inline(always)]
9013        unsafe fn step(
9014            w: *const u8,
9015            x: *const i8,
9016            flip: core::arch::x86_64::__m512i,
9017            acc: core::arch::x86_64::__m512i,
9018        ) -> core::arch::x86_64::__m512i {
9019            unsafe {
9020                use core::arch::x86_64::*;
9021                let wv = _mm512_xor_si512(_mm512_loadu_si512(w as *const _), flip);
9022                _mm512_dpbusd_epi32(acc, wv, _mm512_loadu_si512(x as *const _))
9023            }
9024        }
9025        let (mut a0, mut a1, mut a2, mut a3) = (
9026            _mm512_setzero_si512(),
9027            _mm512_setzero_si512(),
9028            _mm512_setzero_si512(),
9029            _mm512_setzero_si512(),
9030        );
9031        let mut j = 0usize;
9032        while j + 256 <= n {
9033            a0 = step(w.as_ptr().add(j), xq.as_ptr().add(j), flip, a0);
9034            a1 = step(w.as_ptr().add(j + 64), xq.as_ptr().add(j + 64), flip, a1);
9035            a2 = step(w.as_ptr().add(j + 128), xq.as_ptr().add(j + 128), flip, a2);
9036            a3 = step(w.as_ptr().add(j + 192), xq.as_ptr().add(j + 192), flip, a3);
9037            j += 256;
9038        }
9039        while j + 64 <= n {
9040            a0 = step(w.as_ptr().add(j), xq.as_ptr().add(j), flip, a0);
9041            j += 64;
9042        }
9043        let mut total = _mm512_reduce_add_epi32(_mm512_add_epi32(
9044            _mm512_add_epi32(a0, a1),
9045            _mm512_add_epi32(a2, a3),
9046        ));
9047        // Scalar tail: (w as i8) + 128 ≡ (w as u8) ^ 0x80.
9048        while j < n {
9049            total += ((w[j] ^ 0x80) as i32) * xq[j] as i32;
9050            j += 1;
9051        }
9052        total
9053    }
9054}
9055
9056/// One q4 row via AVX2: nibbles → centered i8 (unpacklo/hi restores the
9057/// writer's flat order, same as the NEON vzip pair), maddubs against
9058/// the pre-quantized activation group, × the group's f16 scale. Pair
9059/// saturation safe: |w|≤8, |x|≤127 → 2·8·127 ≪ 32767. Mirror of
9060/// `dot_q4_row_sdot`.
9061#[cfg(target_arch = "x86_64")]
9062#[target_feature(enable = "avx2")]
9063unsafe fn dot_q4_row_avx2(packed: &[u8], scales: &[u8], g0: usize, gpr: usize, xq: &[i8]) -> f32 {
9064    // SAFETY: callers uphold slice-length contracts (16 packed bytes and
9065    // 2 scale bytes per group; xq.len() == gpr·GROUP_SIZE).
9066    unsafe {
9067        use core::arch::x86_64::*;
9068        let lomask = _mm_set1_epi8(0x0F);
9069        let eight = _mm256_set1_epi8(8);
9070        let ones = _mm256_set1_epi16(1);
9071        let mut acc = 0f32;
9072        for gi in 0..gpr {
9073            let g = g0 + gi;
9074            let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
9075            let b = _mm_loadu_si128(packed.as_ptr().add(g * 16) as *const __m128i);
9076            let lo = _mm_and_si128(b, lomask);
9077            let hi = _mm_and_si128(_mm_srli_epi16::<4>(b), lomask);
9078            let w = _mm256_sub_epi8(
9079                _mm256_set_m128i(_mm_unpackhi_epi8(lo, hi), _mm_unpacklo_epi8(lo, hi)),
9080                eight,
9081            );
9082            let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
9083            let p16 = _mm256_maddubs_epi16(_mm256_abs_epi8(w), _mm256_sign_epi8(x, w));
9084            let d = _mm256_madd_epi16(p16, ones);
9085            let hi128 = _mm256_extracti128_si256::<1>(d);
9086            let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
9087            let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
9088            let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
9089            acc += _mm_cvtsi128_si32(s32) as f32 * s;
9090        }
9091        acc
9092    }
9093}
9094
9095/// Two-activation q4 row via AVX2: nibbles unpacked ONCE per group,
9096/// both activations dotted against the same centered i8 register.
9097#[cfg(target_arch = "x86_64")]
9098#[target_feature(enable = "avx2")]
9099unsafe fn dot_q4_row_avx2_2(
9100    packed: &[u8],
9101    scales: &[u8],
9102    g0: usize,
9103    gpr: usize,
9104    xq1: &[i8],
9105    xq2: &[i8],
9106) -> (f32, f32) {
9107    // SAFETY: callers uphold slice-length contracts (see dot_q4_row_avx2).
9108    unsafe {
9109        use core::arch::x86_64::*;
9110        let lomask = _mm_set1_epi8(0x0F);
9111        let eight = _mm256_set1_epi8(8);
9112        let ones = _mm256_set1_epi16(1);
9113        let (mut acc1, mut acc2) = (0f32, 0f32);
9114        #[inline(always)]
9115        unsafe fn hsum(d: core::arch::x86_64::__m256i) -> i32 {
9116            unsafe {
9117                use core::arch::x86_64::*;
9118                let hi128 = _mm256_extracti128_si256::<1>(d);
9119                let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
9120                let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
9121                let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
9122                _mm_cvtsi128_si32(s32)
9123            }
9124        }
9125        for gi in 0..gpr {
9126            let g = g0 + gi;
9127            let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
9128            let b = _mm_loadu_si128(packed.as_ptr().add(g * 16) as *const __m128i);
9129            let lo = _mm_and_si128(b, lomask);
9130            let hi = _mm_and_si128(_mm_srli_epi16::<4>(b), lomask);
9131            let w = _mm256_sub_epi8(
9132                _mm256_set_m128i(_mm_unpackhi_epi8(lo, hi), _mm_unpacklo_epi8(lo, hi)),
9133                eight,
9134            );
9135            let aw = _mm256_abs_epi8(w);
9136            let x1 = _mm256_loadu_si256(xq1.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
9137            let x2 = _mm256_loadu_si256(xq2.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
9138            let d1 = _mm256_madd_epi16(_mm256_maddubs_epi16(aw, _mm256_sign_epi8(x1, w)), ones);
9139            let d2 = _mm256_madd_epi16(_mm256_maddubs_epi16(aw, _mm256_sign_epi8(x2, w)), ones);
9140            acc1 += hsum(d1) as f32 * s;
9141            acc2 += hsum(d2) as f32 * s;
9142        }
9143        (acc1, acc2)
9144    }
9145}
9146
9147/// One q8 row range via AVX2 (x86 mirror of `q8_range_sdot`).
9148#[cfg(target_arch = "x86_64")]
9149fn q8_range_avx2(
9150    q: &[u8],
9151    row_scale: &[f32],
9152    act: &SplitAct,
9153    cols: usize,
9154    out_addr: SendMut,
9155    start: usize,
9156    end: usize,
9157) {
9158    for o in start..end {
9159        let v = row_dot_avx2(&q[o * cols..(o + 1) * cols], act) * row_scale[o];
9160        // SAFETY: disjoint row ranges per worker.
9161        unsafe { *out_addr.at(o) = v };
9162    }
9163}
9164
9165/// Two-input q8 row range via AVX2 (x86 mirror of `q8_range2_sdot`).
9166#[cfg(target_arch = "x86_64")]
9167#[allow(clippy::too_many_arguments)]
9168fn q8_range2_avx2(
9169    q: &[u8],
9170    row_scale: &[f32],
9171    a1: &SplitAct,
9172    a2: &SplitAct,
9173    cols: usize,
9174    p1: SendMut,
9175    p2: SendMut,
9176    start: usize,
9177    end: usize,
9178) {
9179    for o in start..end {
9180        let row = &q[o * cols..(o + 1) * cols];
9181        // SAFETY: disjoint row ranges per worker.
9182        unsafe {
9183            *p1.at(o) = row_dot_avx2(row, a1) * row_scale[o];
9184            *p2.at(o) = row_dot_avx2(row, a2) * row_scale[o];
9185        }
9186    }
9187}
9188
9189// ───────────────────── A8W8 SDOT path (port of vmfcore, ×1.78 decode) ─────────────────────
9190
9191/// ARMv8.6 i8mm (smmla): 32 int8 MACs per instruction vs sdot's 16 —
9192/// yet MEASURED 2.4× SLOWER than the blocked sdot on Apple silicon
9193/// (108 vs 264 GF/s): the on-the-fly vcombine packing and the two-
9194/// accumulator dependency chain swamp the MAC advantage, and Apple's
9195/// four SIMD pipes already keep sdot fed. OPT-IN (CMF_I8MM=1) for
9196/// field trials on Cortex-A710/X-class parts with two pipes, where the
9197/// balance may differ; a pre-interleaved weight layout (repack infra)
9198/// is the known path if it ever earns its keep.
9199#[cfg(target_arch = "aarch64")]
9200fn i8mm_enabled() -> bool {
9201    static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
9202    *ON.get_or_init(|| {
9203        std::env::var("CMF_I8MM").map(|v| v == "1").unwrap_or(false)
9204            && std::arch::is_aarch64_feature_detected!("i8mm")
9205    })
9206}
9207
9208/// SDOT enabled? Default ON when the CPU has ARMv8.2 dotprod;
9209/// `CMF_SDOT=0` disables (falls back to i8×f32 NEON).
9210/// (On non-ARM release builds only the test tolerance switch calls it.)
9211#[cfg_attr(not(target_arch = "aarch64"), allow(dead_code))]
9212fn sdot_enabled() -> bool {
9213    use std::sync::OnceLock;
9214    static ON: OnceLock<bool> = OnceLock::new();
9215    *ON.get_or_init(|| {
9216        let want = std::env::var("CMF_SDOT").map(|v| v != "0").unwrap_or(true);
9217        if !want {
9218            return false;
9219        }
9220
9221        #[cfg(target_arch = "aarch64")]
9222        {
9223            if std::arch::is_aarch64_feature_detected!("dotprod") {
9224                return true;
9225            }
9226            #[cfg(target_os = "android")]
9227            {
9228                if let Ok(cpuinfo) = std::fs::read_to_string("/proc/cpuinfo") {
9229                    if cpuinfo.lines().any(|l| {
9230                        (l.starts_with("Features") || l.starts_with("features"))
9231                            && l.contains("asimddp")
9232                    }) {
9233                        return true;
9234                    }
9235                }
9236            }
9237            false
9238        }
9239        #[cfg(not(target_arch = "aarch64"))]
9240        {
9241            false
9242        }
9243    })
9244}
9245
9246/// Two-field activation split (≡ vmfcore `q8_split_prep`): outlier
9247/// channels (>8·rms) are computed exactly in f32; the bulk (outliers
9248/// zeroed → clean absmax) goes through int8 SDOT. Computed ONCE per
9249/// matvec, shared by all rows/workers.
9250struct SplitAct {
9251    xq: Vec<i8>,
9252    sx: f32,
9253    outliers: Vec<(usize, f32)>,
9254    /// Σ xq — the VNNI bias-trick correction (`(w+128)·x` sums need
9255    /// `−128·Σx`); one i32 per split, computed once per matvec.
9256    #[cfg_attr(not(target_arch = "x86_64"), allow(dead_code))]
9257    xsum: i32,
9258}
9259
9260thread_local! {
9261    /// Recycled xq buffers: split_act runs for every matvec (~200/token)
9262    /// and its hidden-size allocation was steady-state heap churn.
9263    static XQ_FREE: std::cell::RefCell<Vec<Vec<i8>>> =
9264        const { std::cell::RefCell::new(Vec::new()) };
9265}
9266
9267impl Drop for SplitAct {
9268    fn drop(&mut self) {
9269        let buf = std::mem::take(&mut self.xq);
9270        if buf.capacity() > 0 {
9271            XQ_FREE.with(|f| {
9272                let mut f = f.borrow_mut();
9273                if f.len() < 16 {
9274                    f.push(buf);
9275                }
9276            });
9277        }
9278    }
9279}
9280
9281thread_local! {
9282    /// One scratch row per WORKER, kept for the life of the thread.
9283    ///
9284    /// The kernels take a row of group scales per dispatch, and a fresh
9285    /// `vec![0f32; gpr]` inside the closure is one allocation per worker per
9286    /// dispatch — on the release checkpoint about six thousand a token, a
9287    /// quarter of everything the benchmark counts.
9288    static KROW: std::cell::RefCell<Vec<f32>> = const { std::cell::RefCell::new(Vec::new()) };
9289}
9290
9291/// Borrow `n` floats of the calling worker's scratch. Nothing inside a
9292/// kernel body borrows it again, which is what keeps the RefCell honest.
9293#[inline]
9294fn with_krow<R>(n: usize, f: impl FnOnce(&mut [f32]) -> R) -> R {
9295    KROW.with(|s| {
9296        let mut b = s.borrow_mut();
9297        if b.len() < n {
9298            b.resize(n, 0.0);
9299        }
9300        f(&mut b[..n])
9301    })
9302}
9303
9304fn split_act(x: &[f32]) -> SplitAct {
9305    let n = x.len();
9306    let rms = (x.iter().map(|&v| (v * v) as f64).sum::<f64>() / n.max(1) as f64).sqrt() as f32;
9307    let thr = 8.0 * rms;
9308    // One pass: collect outliers and the bulk absmax (outliers excluded —
9309    // identical to the old zero-then-fold over a copied buffer, minus the
9310    // full-vector copy).
9311    let mut outliers: Vec<(usize, f32)> = Vec::new();
9312    let mut amax = 0f32;
9313    for (j, &v) in x.iter().enumerate() {
9314        let a = v.abs();
9315        if a > thr {
9316            outliers.push((j, v));
9317        } else if a > amax {
9318            amax = a;
9319        }
9320    }
9321    let sx = if amax > 0.0 { amax / 127.0 } else { 1.0 };
9322    let inv = 1.0 / sx;
9323    let mut xq = XQ_FREE.with(|f| f.borrow_mut().pop()).unwrap_or_default();
9324    xq.clear();
9325    xq.reserve(n);
9326    if outliers.is_empty() {
9327        xq.extend(
9328            x.iter()
9329                .map(|&v| (v * inv).round().clamp(-127.0, 127.0) as i8),
9330        );
9331    } else {
9332        // Outlier slots quantize to 0 (their exact term is added later).
9333        xq.extend(x.iter().map(|&v| {
9334            if v.abs() > thr {
9335                0
9336            } else {
9337                (v * inv).round().clamp(-127.0, 127.0) as i8
9338            }
9339        }));
9340    }
9341    let xsum = xq.iter().map(|&v| v as i32).sum();
9342    SplitAct {
9343        xq,
9344        sx,
9345        outliers,
9346        xsum,
9347    }
9348}
9349
9350fn split_act_q8_2f(x: &[f32], col: &[f32]) -> SplitAct {
9351    let n = x.len();
9352    let rms = (x
9353        .iter()
9354        .zip(col)
9355        .map(|(&a, &c)| {
9356            let v = a * c;
9357            (v * v) as f64
9358        })
9359        .sum::<f64>()
9360        / n.max(1) as f64)
9361        .sqrt() as f32;
9362    let thr = 8.0 * rms;
9363
9364    let mut outliers = Vec::new();
9365    let mut amax = 0f32;
9366    for (j, (&a, &c)) in x.iter().zip(col).enumerate() {
9367        let v = a * c;
9368        let s = v.abs();
9369        if s > thr {
9370            outliers.push((j, v));
9371        } else if s > amax {
9372            amax = s;
9373        }
9374    }
9375
9376    let sx = if amax > 0.0 { amax / 127.0 } else { 1.0 };
9377    let inv = 1.0 / sx;
9378    let mut xq = XQ_FREE.with(|f| f.borrow_mut().pop()).unwrap_or_default();
9379    xq.clear();
9380    xq.reserve(n);
9381    if outliers.is_empty() {
9382        xq.extend(
9383            x.iter()
9384                .zip(col)
9385                .map(|(&a, &c)| ((a * c) * inv).round().clamp(-127.0, 127.0) as i8),
9386        );
9387    } else {
9388        xq.extend(x.iter().zip(col).map(|(&a, &c)| {
9389            let v = a * c;
9390            if v.abs() > thr {
9391                0
9392            } else {
9393                (v * inv).round().clamp(-127.0, 127.0) as i8
9394            }
9395        }));
9396    }
9397    let xsum = xq.iter().map(|&v| v as i32).sum();
9398    SplitAct {
9399        xq,
9400        sx,
9401        outliers,
9402        xsum,
9403    }
9404}
9405
9406/// int8(weight)·int8(activation) → i32 via `sdot` (inline asm — the
9407/// vdotq intrinsic is unstable; port of vmfcore `dot_i8_sdot`).
9408#[cfg(target_arch = "aarch64")]
9409#[target_feature(enable = "neon,dotprod")]
9410unsafe fn dot_i8_sdot(w: &[u8], xq: &[i8]) -> i32 {
9411    // SAFETY: callers uphold slice-length contracts (see call sites).
9412    unsafe {
9413        use core::arch::aarch64::*;
9414        use core::arch::asm;
9415        let wp = w.as_ptr() as *const i8;
9416        let n = w.len();
9417        let (mut a0, mut a1, mut a2, mut a3) = (
9418            vdupq_n_s32(0),
9419            vdupq_n_s32(0),
9420            vdupq_n_s32(0),
9421            vdupq_n_s32(0),
9422        );
9423        let mut i = 0;
9424        while i + 64 <= n {
9425            let (w0, x0) = (vld1q_s8(wp.add(i)), vld1q_s8(xq.as_ptr().add(i)));
9426            let (w1, x1) = (vld1q_s8(wp.add(i + 16)), vld1q_s8(xq.as_ptr().add(i + 16)));
9427            let (w2, x2) = (vld1q_s8(wp.add(i + 32)), vld1q_s8(xq.as_ptr().add(i + 32)));
9428            let (w3, x3) = (vld1q_s8(wp.add(i + 48)), vld1q_s8(xq.as_ptr().add(i + 48)));
9429            asm!(
9430                "sdot {a0:v}.4s, {w0:v}.16b, {x0:v}.16b",
9431                "sdot {a1:v}.4s, {w1:v}.16b, {x1:v}.16b",
9432                "sdot {a2:v}.4s, {w2:v}.16b, {x2:v}.16b",
9433                "sdot {a3:v}.4s, {w3:v}.16b, {x3:v}.16b",
9434                a0 = inout(vreg) a0, a1 = inout(vreg) a1, a2 = inout(vreg) a2, a3 = inout(vreg) a3,
9435                w0 = in(vreg) w0, x0 = in(vreg) x0, w1 = in(vreg) w1, x1 = in(vreg) x1,
9436                w2 = in(vreg) w2, x2 = in(vreg) x2, w3 = in(vreg) w3, x3 = in(vreg) x3,
9437                options(pure, nomem, nostack),
9438            );
9439            i += 64;
9440        }
9441        while i + 16 <= n {
9442            let (wv, xv) = (vld1q_s8(wp.add(i)), vld1q_s8(xq.as_ptr().add(i)));
9443            asm!("sdot {a:v}.4s, {w:v}.16b, {x:v}.16b",
9444                 a = inout(vreg) a0, w = in(vreg) wv, x = in(vreg) xv, options(pure, nomem, nostack));
9445            i += 16;
9446        }
9447        let mut s = vaddvq_s32(vaddq_s32(vaddq_s32(a0, a1), vaddq_s32(a2, a3)));
9448        while i < n {
9449            s += (*wp.add(i)) as i32 * xq[i] as i32;
9450            i += 1;
9451        }
9452        s
9453    }
9454}
9455
9456/// Row-blocked SDOT: 4 output rows per pass — the activation chunk is
9457/// loaded once and reused, 4 independent accumulators hide sdot latency
9458/// (port of vmfcore `dot_i8_sdot_4rows`).
9459#[cfg(target_arch = "aarch64")]
9460#[target_feature(enable = "neon,dotprod")]
9461unsafe fn dot_i8_sdot_4rows(w0: &[u8], w1: &[u8], w2: &[u8], w3: &[u8], xq: &[i8]) -> [i32; 4] {
9462    // SAFETY: callers uphold slice-length contracts (see call sites).
9463    unsafe {
9464        use core::arch::aarch64::*;
9465        use core::arch::asm;
9466        let n = xq.len();
9467        let px = xq.as_ptr();
9468        let (p0, p1, p2, p3) = (
9469            w0.as_ptr() as *const i8,
9470            w1.as_ptr() as *const i8,
9471            w2.as_ptr() as *const i8,
9472            w3.as_ptr() as *const i8,
9473        );
9474        let (mut a0, mut a1, mut a2, mut a3) = (
9475            vdupq_n_s32(0),
9476            vdupq_n_s32(0),
9477            vdupq_n_s32(0),
9478            vdupq_n_s32(0),
9479        );
9480        let mut i = 0;
9481        while i + 16 <= n {
9482            let x = vld1q_s8(px.add(i));
9483            let v0 = vld1q_s8(p0.add(i));
9484            let v1 = vld1q_s8(p1.add(i));
9485            let v2 = vld1q_s8(p2.add(i));
9486            let v3 = vld1q_s8(p3.add(i));
9487            asm!(
9488                "sdot {a0:v}.4s, {v0:v}.16b, {x:v}.16b",
9489                "sdot {a1:v}.4s, {v1:v}.16b, {x:v}.16b",
9490                "sdot {a2:v}.4s, {v2:v}.16b, {x:v}.16b",
9491                "sdot {a3:v}.4s, {v3:v}.16b, {x:v}.16b",
9492                a0 = inout(vreg) a0, a1 = inout(vreg) a1, a2 = inout(vreg) a2, a3 = inout(vreg) a3,
9493                v0 = in(vreg) v0, v1 = in(vreg) v1, v2 = in(vreg) v2, v3 = in(vreg) v3, x = in(vreg) x,
9494                options(pure, nomem, nostack),
9495            );
9496            i += 16;
9497        }
9498        let mut r = [
9499            vaddvq_s32(a0),
9500            vaddvq_s32(a1),
9501            vaddvq_s32(a2),
9502            vaddvq_s32(a3),
9503        ];
9504        while i < n {
9505            let xi = *px.add(i) as i32;
9506            r[0] += (*p0.add(i)) as i32 * xi;
9507            r[1] += (*p1.add(i)) as i32 * xi;
9508            r[2] += (*p2.add(i)) as i32 * xi;
9509            r[3] += (*p3.add(i)) as i32 * xi;
9510            i += 1;
9511        }
9512        r
9513    }
9514}
9515
9516/// 4 interleaved rows in one pass: the repacked group is [r0[c], r1[c],
9517/// r2[c], r3[c]] per 16-byte chunk, so each iteration reads ONE 64-byte
9518/// line plus the shared activation chunk — a single sequential weight
9519/// stream per worker. Per-row accumulation is the same one-accumulator
9520/// scheme as `dot_i8_sdot_4rows`; integer sums are exact, so outputs
9521/// are bit-identical to the mmap-layout kernel.
9522#[cfg(target_arch = "aarch64")]
9523#[target_feature(enable = "neon,dotprod")]
9524unsafe fn dot_i8_sdot_4rows_il(g: &[u8], xq: &[i8]) -> [i32; 4] {
9525    // SAFETY: callers uphold slice-length contracts (g.len() == 4·n,
9526    // n % 16 == 0 — guaranteed by the repack gate).
9527    unsafe {
9528        use core::arch::aarch64::*;
9529        use core::arch::asm;
9530        let n = xq.len();
9531        let px = xq.as_ptr();
9532        let pg = g.as_ptr() as *const i8;
9533        let (mut a0, mut a1, mut a2, mut a3) = (
9534            vdupq_n_s32(0),
9535            vdupq_n_s32(0),
9536            vdupq_n_s32(0),
9537            vdupq_n_s32(0),
9538        );
9539        let mut i = 0;
9540        while i + 16 <= n {
9541            let x = vld1q_s8(px.add(i));
9542            let base = pg.add(4 * i);
9543            let v0 = vld1q_s8(base);
9544            let v1 = vld1q_s8(base.add(16));
9545            let v2 = vld1q_s8(base.add(32));
9546            let v3 = vld1q_s8(base.add(48));
9547            asm!(
9548                "sdot {a0:v}.4s, {v0:v}.16b, {x:v}.16b",
9549                "sdot {a1:v}.4s, {v1:v}.16b, {x:v}.16b",
9550                "sdot {a2:v}.4s, {v2:v}.16b, {x:v}.16b",
9551                "sdot {a3:v}.4s, {v3:v}.16b, {x:v}.16b",
9552                a0 = inout(vreg) a0, a1 = inout(vreg) a1, a2 = inout(vreg) a2, a3 = inout(vreg) a3,
9553                v0 = in(vreg) v0, v1 = in(vreg) v1, v2 = in(vreg) v2, v3 = in(vreg) v3, x = in(vreg) x,
9554                options(pure, nomem, nostack),
9555            );
9556            i += 16;
9557        }
9558        [
9559            vaddvq_s32(a0),
9560            vaddvq_s32(a1),
9561            vaddvq_s32(a2),
9562            vaddvq_s32(a3),
9563        ]
9564    }
9565}
9566
9567/// One q8 row range via SDOT (4-row blocks + tail) — the body of
9568/// `qmatvec`'s hot loop, extracted so multi-matrix jobs can drive the
9569/// SAME kernel for several tensors under one pool dispatch. `rep` — the
9570/// load-time interleaved repack (empty = mmap layout only); rows outside
9571/// full 4-row groups always come from the mmap layout.
9572#[cfg(target_arch = "aarch64")]
9573fn q8_range_sdot(
9574    q: &[u8],
9575    rep: &[u8],
9576    row_scale: &[f32],
9577    act: &SplitAct,
9578    cols: usize,
9579    out_addr: SendMut,
9580    start: usize,
9581    end: usize,
9582) {
9583    let mut o = start;
9584    // Leading rows to the group boundary (repack path only): the pool
9585    // splits row ranges arbitrarily, groups are absolute.
9586    if !rep.is_empty() {
9587        while o < end && o % 4 != 0 {
9588            let v = row_dot_sdot(&q[o * cols..(o + 1) * cols], act) * row_scale[o];
9589            unsafe { *out_addr.at(o) = v };
9590            o += 1;
9591        }
9592    }
9593    while o + 4 <= end {
9594        let r = if rep.is_empty() {
9595            unsafe {
9596                dot_i8_sdot_4rows(
9597                    &q[o * cols..(o + 1) * cols],
9598                    &q[(o + 1) * cols..(o + 2) * cols],
9599                    &q[(o + 2) * cols..(o + 3) * cols],
9600                    &q[(o + 3) * cols..(o + 4) * cols],
9601                    &act.xq,
9602                )
9603            }
9604        } else {
9605            unsafe { dot_i8_sdot_4rows_il(&rep[o * cols..(o + 4) * cols], &act.xq) }
9606        };
9607        for k in 0..4 {
9608            let mut acc = r[k] as f32 * act.sx;
9609            for &(j, xv) in &act.outliers {
9610                acc += (q[(o + k) * cols + j] as i8) as f32 * xv;
9611            }
9612            // SAFETY: disjoint row ranges per worker.
9613            unsafe { *out_addr.at(o + k) = acc * row_scale[o + k] };
9614        }
9615        o += 4;
9616    }
9617    while o < end {
9618        let v = row_dot_sdot(&q[o * cols..(o + 1) * cols], act) * row_scale[o];
9619        unsafe { *out_addr.at(o) = v };
9620        o += 1;
9621    }
9622}
9623
9624/// Two-input q8 row range via SDOT — `qmatvec2`'s hot loop, extracted
9625/// for the fused pair multi-matrix job (`matvec2_many`).
9626#[cfg(target_arch = "aarch64")]
9627#[allow(clippy::too_many_arguments)]
9628fn q8_range2_sdot(
9629    q: &[u8],
9630    row_scale: &[f32],
9631    a1: &SplitAct,
9632    a2: &SplitAct,
9633    cols: usize,
9634    p1: SendMut,
9635    p2: SendMut,
9636    start: usize,
9637    end: usize,
9638) {
9639    for o in start..end {
9640        let row = &q[o * cols..(o + 1) * cols];
9641        // SAFETY: disjoint row ranges per worker.
9642        unsafe {
9643            *p1.at(o) = row_dot_sdot(row, a1) * row_scale[o];
9644            *p2.at(o) = row_dot_sdot(row, a2) * row_scale[o];
9645        }
9646    }
9647}
9648
9649/// Two-input q8 row range, f32 kernel (non-SDOT) — same extraction.
9650#[allow(clippy::too_many_arguments)]
9651fn q8_range2_f32(
9652    q: &[u8],
9653    row_scale: &[f32],
9654    x1: &[f32],
9655    x2: &[f32],
9656    cols: usize,
9657    p1: SendMut,
9658    p2: SendMut,
9659    start: usize,
9660    end: usize,
9661) {
9662    for o in start..end {
9663        let row = &q[o * cols..(o + 1) * cols];
9664        // SAFETY: disjoint row ranges per worker.
9665        unsafe {
9666            *p1.at(o) = dot_i8_f32(row, x1) * row_scale[o];
9667            *p2.at(o) = dot_i8_f32(row, x2) * row_scale[o];
9668        }
9669    }
9670}
9671
9672/// Scalar/NEON-f32 q8 row range (non-SDOT platforms) — same extraction.
9673fn q8_range_f32(
9674    q: &[u8],
9675    row_scale: &[f32],
9676    xs: &[f32],
9677    cols: usize,
9678    out_addr: SendMut,
9679    start: usize,
9680    end: usize,
9681) {
9682    for o in start..end {
9683        let v = dot_i8_f32(&q[o * cols..(o + 1) * cols], xs) * row_scale[o];
9684        // SAFETY: disjoint row ranges per worker.
9685        unsafe { *out_addr.at(o) = v };
9686    }
9687}
9688
9689/// One q8 row against a split activation, portable: the per-arch fast
9690/// dots where they exist, the exact scalar loop elsewhere. The scalar
9691/// arm is also the test oracle for both fast arms.
9692#[inline]
9693fn q8_row_dot(row: &[u8], act: &SplitAct) -> f32 {
9694    #[cfg(target_arch = "aarch64")]
9695    return row_dot_sdot(row, act);
9696    #[cfg(target_arch = "x86_64")]
9697    return row_dot_avx2(row, act);
9698    #[allow(unreachable_code)]
9699    q8_row_dot_scalar(row, act)
9700}
9701
9702#[allow(dead_code)]
9703fn q8_row_dot_scalar(row: &[u8], act: &SplitAct) -> f32 {
9704    let mut acc = 0i32;
9705    for (k, &b) in row.iter().enumerate() {
9706        acc += (b as i8) as i32 * act.xq[k] as i32;
9707    }
9708    let mut acc = acc as f32 * act.sx;
9709    for &(j, xv) in &act.outliers {
9710        acc += (row[j] as i8) as f32 * xv;
9711    }
9712    acc
9713}
9714
9715/// SDOT row dot with exact outlier correction:
9716/// `dot = sdot(w, xq)·sx + Σ_outl w[j]·x[j]` (then × row_scale by caller).
9717#[cfg(target_arch = "aarch64")]
9718#[inline]
9719fn row_dot_sdot(row: &[u8], act: &SplitAct) -> f32 {
9720    let mut acc = unsafe { dot_i8_sdot(row, &act.xq) } as f32 * act.sx;
9721    for &(j, xv) in &act.outliers {
9722        acc += (row[j] as i8) as f32 * xv;
9723    }
9724    acc
9725}
9726
9727/// One q4 row via SDOT: each 32-group's nibbles unpack to centered i8
9728/// (nib−8 ∈ [−8,7]), int8×int8 `sdot` against the pre-quantized
9729/// activation group, × the group's f16 scale. Returns Σ_g dot_g·s_g;
9730/// the caller multiplies by the activation scale and adds the exact
9731/// outlier terms (port of vmfcore `dot_q4_block_sdot`, +23% measured).
9732/// Nibble order matches the writer: element 2k = low nibble, 2k+1 = high
9733/// → zip(lo,hi) restores flat order.
9734#[cfg(target_arch = "aarch64")]
9735#[target_feature(enable = "neon,dotprod")]
9736unsafe fn dot_q4_row_sdot(packed: &[u8], scales: &[u8], g0: usize, gpr: usize, xq: &[i8]) -> f32 {
9737    // SAFETY: callers uphold slice-length contracts (16 packed bytes and
9738    // 2 scale bytes per group; xq.len() == gpr·GROUP_SIZE).
9739    unsafe {
9740        use core::arch::aarch64::*;
9741        use core::arch::asm;
9742        let lomask = vdupq_n_u8(0x0F);
9743        let eight = vdupq_n_s8(8);
9744        let mut acc = 0f32;
9745        for gi in 0..gpr {
9746            let g = g0 + gi;
9747            let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
9748            let b = vld1q_u8(packed.as_ptr().add(g * 16));
9749            let lo = vandq_u8(b, lomask);
9750            let hi = vshrq_n_u8::<4>(b);
9751            let e0 = vsubq_s8(vreinterpretq_s8_u8(vzip1q_u8(lo, hi)), eight);
9752            let e1 = vsubq_s8(vreinterpretq_s8_u8(vzip2q_u8(lo, hi)), eight);
9753            let x0 = vld1q_s8(xq.as_ptr().add(gi * GROUP_SIZE));
9754            let x1 = vld1q_s8(xq.as_ptr().add(gi * GROUP_SIZE + 16));
9755            let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
9756            asm!(
9757                "sdot {a0:v}.4s, {e0:v}.16b, {x0:v}.16b",
9758                "sdot {a1:v}.4s, {e1:v}.16b, {x1:v}.16b",
9759                a0 = inout(vreg) a0, a1 = inout(vreg) a1,
9760                e0 = in(vreg) e0, x0 = in(vreg) x0, e1 = in(vreg) e1, x1 = in(vreg) x1,
9761                options(pure, nomem, nostack),
9762            );
9763            acc += vaddvq_s32(vaddq_s32(a0, a1)) as f32 * s;
9764        }
9765        acc
9766    }
9767}
9768
9769/// Two-activation q4 row via SDOT: the nibble unpack (the expensive
9770/// part) happens ONCE per group; both pre-quantized activations are
9771/// dotted against the same centered i8 registers. Per-lane math matches
9772/// `dot_q4_row_sdot` exactly.
9773#[cfg(target_arch = "aarch64")]
9774#[target_feature(enable = "neon,dotprod")]
9775unsafe fn dot_q4_row_sdot2(
9776    packed: &[u8],
9777    scales: &[u8],
9778    g0: usize,
9779    gpr: usize,
9780    xq1: &[i8],
9781    xq2: &[i8],
9782) -> (f32, f32) {
9783    // SAFETY: callers uphold slice-length contracts (16 packed bytes and
9784    // 2 scale bytes per group; xq*.len() == gpr·GROUP_SIZE).
9785    unsafe {
9786        use core::arch::aarch64::*;
9787        use core::arch::asm;
9788        let lomask = vdupq_n_u8(0x0F);
9789        let eight = vdupq_n_s8(8);
9790        let (mut acc1, mut acc2) = (0f32, 0f32);
9791        for gi in 0..gpr {
9792            let g = g0 + gi;
9793            let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
9794            let b = vld1q_u8(packed.as_ptr().add(g * 16));
9795            let lo = vandq_u8(b, lomask);
9796            let hi = vshrq_n_u8::<4>(b);
9797            let e0 = vsubq_s8(vreinterpretq_s8_u8(vzip1q_u8(lo, hi)), eight);
9798            let e1 = vsubq_s8(vreinterpretq_s8_u8(vzip2q_u8(lo, hi)), eight);
9799            let x10 = vld1q_s8(xq1.as_ptr().add(gi * GROUP_SIZE));
9800            let x11 = vld1q_s8(xq1.as_ptr().add(gi * GROUP_SIZE + 16));
9801            let x20 = vld1q_s8(xq2.as_ptr().add(gi * GROUP_SIZE));
9802            let x21 = vld1q_s8(xq2.as_ptr().add(gi * GROUP_SIZE + 16));
9803            let (mut a0, mut a1, mut b0, mut b1) = (
9804                vdupq_n_s32(0),
9805                vdupq_n_s32(0),
9806                vdupq_n_s32(0),
9807                vdupq_n_s32(0),
9808            );
9809            asm!(
9810                "sdot {a0:v}.4s, {e0:v}.16b, {x10:v}.16b",
9811                "sdot {a1:v}.4s, {e1:v}.16b, {x11:v}.16b",
9812                "sdot {b0:v}.4s, {e0:v}.16b, {x20:v}.16b",
9813                "sdot {b1:v}.4s, {e1:v}.16b, {x21:v}.16b",
9814                a0 = inout(vreg) a0, a1 = inout(vreg) a1,
9815                b0 = inout(vreg) b0, b1 = inout(vreg) b1,
9816                e0 = in(vreg) e0, e1 = in(vreg) e1,
9817                x10 = in(vreg) x10, x11 = in(vreg) x11,
9818                x20 = in(vreg) x20, x21 = in(vreg) x21,
9819                options(pure, nomem, nostack),
9820            );
9821            acc1 += vaddvq_s32(vaddq_s32(a0, a1)) as f32 * s;
9822            acc2 += vaddvq_s32(vaddq_s32(b0, b1)) as f32 * s;
9823        }
9824        (acc1, acc2)
9825    }
9826}
9827
9828// ───────────────────── fused int8 kernels ─────────────────────
9829
9830/// `acc += w · row` where the row is centered i8 — NEON widen+fma on
9831/// aarch64, scalar elsewhere. The KV-cache q8 value path rides on this.
9832#[inline]
9833pub(crate) fn axpy_i8_f32(acc: &mut [f32], row: &[i8], w: f32) {
9834    #[cfg(target_arch = "aarch64")]
9835    unsafe {
9836        return axpy_i8_f32_neon(acc, row, w);
9837    }
9838    #[cfg(target_arch = "x86_64")]
9839    if avx2_enabled() {
9840        return unsafe { axpy_i8_f32_avx2(acc, row, w) };
9841    }
9842    #[allow(unreachable_code)]
9843    {
9844        for (a, &b) in acc.iter_mut().zip(row) {
9845            *a += w * b as f32;
9846        }
9847    }
9848}
9849
9850/// i8→f32 axpy via AVX2/FMA (x86 mirror of `axpy_i8_f32_neon`).
9851#[cfg(target_arch = "x86_64")]
9852#[target_feature(enable = "avx2,fma")]
9853unsafe fn axpy_i8_f32_avx2(acc: &mut [f32], row: &[i8], w: f32) {
9854    // SAFETY: callers uphold slice-length contracts (see call sites).
9855    unsafe {
9856        use core::arch::x86_64::*;
9857        let n = acc.len().min(row.len());
9858        let ap = acc.as_mut_ptr();
9859        let rp = row.as_ptr();
9860        let wv = _mm256_set1_ps(w);
9861        let mut j = 0usize;
9862        while j + 16 <= n {
9863            let rb = _mm_loadu_si128(rp.add(j) as *const __m128i);
9864            let lo = _mm256_cvtepi32_ps(_mm256_cvtepi8_epi32(rb));
9865            let hi = _mm256_cvtepi32_ps(_mm256_cvtepi8_epi32(_mm_srli_si128::<8>(rb)));
9866            let v0 = _mm256_fmadd_ps(wv, lo, _mm256_loadu_ps(ap.add(j)));
9867            let v1 = _mm256_fmadd_ps(wv, hi, _mm256_loadu_ps(ap.add(j + 8)));
9868            _mm256_storeu_ps(ap.add(j), v0);
9869            _mm256_storeu_ps(ap.add(j + 8), v1);
9870            j += 16;
9871        }
9872        while j < n {
9873            *ap.add(j) += w * (*rp.add(j)) as f32;
9874            j += 1;
9875        }
9876    }
9877}
9878
9879#[cfg(target_arch = "aarch64")]
9880#[target_feature(enable = "neon")]
9881unsafe fn axpy_i8_f32_neon(acc: &mut [f32], row: &[i8], w: f32) {
9882    // SAFETY: callers uphold slice-length contracts (see call sites).
9883    unsafe {
9884        use core::arch::aarch64::*;
9885        let n = acc.len().min(row.len());
9886        let ap = acc.as_mut_ptr();
9887        let rp = row.as_ptr();
9888        let wv = vdupq_n_f32(w);
9889        let mut j = 0usize;
9890        while j + 16 <= n {
9891            let rb = vld1q_s8(rp.add(j));
9892            let lo = vmovl_s8(vget_low_s8(rb));
9893            let hi = vmovl_s8(vget_high_s8(rb));
9894            for (off, half) in [(0, lo), (8, hi)] {
9895                let f0 = vcvtq_f32_s32(vmovl_s16(vget_low_s16(half)));
9896                let f1 = vcvtq_f32_s32(vmovl_s16(vget_high_s16(half)));
9897                let o = j + off;
9898                vst1q_f32(ap.add(o), vfmaq_f32(vld1q_f32(ap.add(o)), wv, f0));
9899                vst1q_f32(ap.add(o + 4), vfmaq_f32(vld1q_f32(ap.add(o + 4)), wv, f1));
9900            }
9901            j += 16;
9902        }
9903        while j < n {
9904            *ap.add(j) += w * (*rp.add(j)) as f32;
9905            j += 1;
9906        }
9907    }
9908}
9909
9910/// i8 row · f32 x. NEON on aarch64 (ported from vmfcore `dot_i8_f32_neon`,
9911/// ≈9× scalar), scalar elsewhere.
9912#[inline]
9913pub(crate) fn dot_i8_f32(w: &[u8], x: &[f32]) -> f32 {
9914    #[cfg(target_arch = "aarch64")]
9915    unsafe {
9916        return dot_i8_f32_neon(w, x);
9917    }
9918    #[cfg(target_arch = "x86_64")]
9919    if avx2_enabled() {
9920        return unsafe { dot_i8_f32_avx2(w, x) };
9921    }
9922    #[allow(unreachable_code)]
9923    {
9924        let mut sum = 0.0f32;
9925        for (j, &b) in w.iter().enumerate() {
9926            sum += (b as i8) as f32 * x[j];
9927        }
9928        sum
9929    }
9930}
9931
9932/// i8 row · (x ⊙ col_field) — the q8_2f row dot with the θ col-field
9933/// folded into the product (no prescaled copy of x). NEON on aarch64,
9934/// scalar elsewhere. Used by the active-neuron path `row_dot`.
9935#[inline]
9936fn dot_i8_col_f32(w: &[u8], x: &[f32], col: &[f32]) -> f32 {
9937    #[cfg(target_arch = "aarch64")]
9938    unsafe {
9939        return dot_i8_col_f32_neon(w, x, col);
9940    }
9941    #[allow(unreachable_code)]
9942    {
9943        let mut sum = 0.0f32;
9944        for (j, &b) in w.iter().enumerate() {
9945            sum += (b as i8) as f32 * x[j] * col[j];
9946        }
9947        sum
9948    }
9949}
9950
9951#[cfg(target_arch = "aarch64")]
9952#[target_feature(enable = "neon")]
9953unsafe fn dot_i8_col_f32_neon(w: &[u8], x: &[f32], col: &[f32]) -> f32 {
9954    // SAFETY: callers uphold slice-length contracts (see call sites).
9955    unsafe {
9956        use core::arch::aarch64::*;
9957        let n = x.len();
9958        let wp = w.as_ptr() as *const i8;
9959        let xp = x.as_ptr();
9960        let cp = col.as_ptr();
9961        let (mut a0, mut a1, mut a2, mut a3) = (
9962            vdupq_n_f32(0.0),
9963            vdupq_n_f32(0.0),
9964            vdupq_n_f32(0.0),
9965            vdupq_n_f32(0.0),
9966        );
9967        let mut j = 0usize;
9968        while j + 16 <= n {
9969            let wb = vld1q_s8(wp.add(j));
9970            let lo = vmovl_s8(vget_low_s8(wb));
9971            let hi = vmovl_s8(vget_high_s8(wb));
9972            let w0 = vcvtq_f32_s32(vmovl_s16(vget_low_s16(lo)));
9973            let w1 = vcvtq_f32_s32(vmovl_s16(vget_high_s16(lo)));
9974            let w2 = vcvtq_f32_s32(vmovl_s16(vget_low_s16(hi)));
9975            let w3 = vcvtq_f32_s32(vmovl_s16(vget_high_s16(hi)));
9976            a0 = vfmaq_f32(
9977                a0,
9978                w0,
9979                vmulq_f32(vld1q_f32(xp.add(j)), vld1q_f32(cp.add(j))),
9980            );
9981            a1 = vfmaq_f32(
9982                a1,
9983                w1,
9984                vmulq_f32(vld1q_f32(xp.add(j + 4)), vld1q_f32(cp.add(j + 4))),
9985            );
9986            a2 = vfmaq_f32(
9987                a2,
9988                w2,
9989                vmulq_f32(vld1q_f32(xp.add(j + 8)), vld1q_f32(cp.add(j + 8))),
9990            );
9991            a3 = vfmaq_f32(
9992                a3,
9993                w3,
9994                vmulq_f32(vld1q_f32(xp.add(j + 12)), vld1q_f32(cp.add(j + 12))),
9995            );
9996            j += 16;
9997        }
9998        let mut sum = vaddvq_f32(vaddq_f32(vaddq_f32(a0, a1), vaddq_f32(a2, a3)));
9999        while j < n {
10000            sum += (*wp.add(j)) as f32 * *xp.add(j) * *cp.add(j);
10001            j += 1;
10002        }
10003        sum
10004    }
10005}
10006
10007#[cfg(target_arch = "aarch64")]
10008#[target_feature(enable = "neon")]
10009unsafe fn dot_i8_f32_neon(w: &[u8], x: &[f32]) -> f32 {
10010    // SAFETY: callers uphold slice-length contracts (see call sites).
10011    unsafe {
10012        use core::arch::aarch64::*;
10013        let n = x.len();
10014        let wp = w.as_ptr() as *const i8;
10015        let xp = x.as_ptr();
10016        let (mut a0, mut a1, mut a2, mut a3) = (
10017            vdupq_n_f32(0.0),
10018            vdupq_n_f32(0.0),
10019            vdupq_n_f32(0.0),
10020            vdupq_n_f32(0.0),
10021        );
10022        let mut j = 0usize;
10023        while j + 16 <= n {
10024            let wb = vld1q_s8(wp.add(j));
10025            let lo = vmovl_s8(vget_low_s8(wb));
10026            let hi = vmovl_s8(vget_high_s8(wb));
10027            let w0 = vcvtq_f32_s32(vmovl_s16(vget_low_s16(lo)));
10028            let w1 = vcvtq_f32_s32(vmovl_s16(vget_high_s16(lo)));
10029            let w2 = vcvtq_f32_s32(vmovl_s16(vget_low_s16(hi)));
10030            let w3 = vcvtq_f32_s32(vmovl_s16(vget_high_s16(hi)));
10031            a0 = vfmaq_f32(a0, w0, vld1q_f32(xp.add(j)));
10032            a1 = vfmaq_f32(a1, w1, vld1q_f32(xp.add(j + 4)));
10033            a2 = vfmaq_f32(a2, w2, vld1q_f32(xp.add(j + 8)));
10034            a3 = vfmaq_f32(a3, w3, vld1q_f32(xp.add(j + 12)));
10035            j += 16;
10036        }
10037        let mut sum = vaddvq_f32(vaddq_f32(vaddq_f32(a0, a1), vaddq_f32(a2, a3)));
10038        while j < n {
10039            sum += (*wp.add(j)) as f32 * *xp.add(j);
10040            j += 1;
10041        }
10042        sum
10043    }
10044}
10045
10046#[allow(clippy::too_many_arguments)]
10047fn qmatvec(
10048    q: &[u8],
10049    rep: &[u8],
10050    row_scale: &[f32],
10051    x: &[f32],
10052    col_field: &[f32],
10053    dtype: TensorDtype,
10054    rows: usize,
10055    cols: usize,
10056    out: &mut [f32],
10057    pool: Option<&Pool>,
10058) {
10059    debug_assert_eq!(out.len(), rows);
10060    #[cfg(not(target_arch = "aarch64"))]
10061    let _ = rep;
10062
10063    #[cfg(target_arch = "aarch64")]
10064    if sdot_enabled() {
10065        let act = if dtype == TensorDtype::Q8_2f {
10066            split_act_q8_2f(x, col_field)
10067        } else {
10068            split_act(x)
10069        };
10070        let out_addr = SendMut(out.as_mut_ptr());
10071        let run_range = |start: usize, end: usize| {
10072            q8_range_sdot(q, rep, row_scale, &act, cols, out_addr, start, end)
10073        };
10074        match pool {
10075            Some(pool) if rows >= 256 => pool.run_rows(rows, &run_range),
10076            _ => run_range(0, rows),
10077        }
10078        return;
10079    }
10080    // x86 A8W8 via AVX2 maddubs — same quantized-activation contract as
10081    // the SDOT path (CMF_AVX2=0 keeps the exact i8×f32 loop).
10082    #[cfg(target_arch = "x86_64")]
10083    if avx2_a8w8_enabled() {
10084        let act = if dtype == TensorDtype::Q8_2f {
10085            split_act_q8_2f(x, col_field)
10086        } else {
10087            split_act(x)
10088        };
10089        let out_addr = SendMut(out.as_mut_ptr());
10090        let run_range = |start: usize, end: usize| {
10091            q8_range_avx2(q, row_scale, &act, cols, out_addr, start, end)
10092        };
10093        match pool {
10094            Some(pool) if rows >= 256 => pool.run_rows(rows, &run_range),
10095            _ => run_range(0, rows),
10096        }
10097        return;
10098    }
10099
10100    prescale_with(x, col_field, dtype, 1, |xs| {
10101        let out_addr = SendMut(out.as_mut_ptr());
10102        let run_range = move |start: usize, end: usize| {
10103            for o in start..end {
10104                let v = dot_i8_f32(&q[o * cols..(o + 1) * cols], xs) * row_scale[o];
10105                // SAFETY: disjoint row ranges per worker.
10106                unsafe { *out_addr.at(o) = v };
10107            }
10108        };
10109        match pool {
10110            Some(pool) if rows >= 256 => pool.run_rows(rows, &run_range),
10111            _ => run_range(0, rows),
10112        }
10113    });
10114}
10115
10116#[allow(clippy::too_many_arguments)]
10117fn qmatvec2(
10118    q: &[u8],
10119    row_scale: &[f32],
10120    x1: &[f32],
10121    x2: &[f32],
10122    col_field: &[f32],
10123    dtype: TensorDtype,
10124    rows: usize,
10125    cols: usize,
10126    o1: &mut [f32],
10127    o2: &mut [f32],
10128    pool: Option<&Pool>,
10129) {
10130    #[cfg(target_arch = "aarch64")]
10131    if sdot_enabled() {
10132        let a1s = if dtype == TensorDtype::Q8_2f {
10133            split_act_q8_2f(x1, col_field)
10134        } else {
10135            split_act(x1)
10136        };
10137        let a2s = if dtype == TensorDtype::Q8_2f {
10138            split_act_q8_2f(x2, col_field)
10139        } else {
10140            split_act(x2)
10141        };
10142        let p1 = SendMut(o1.as_mut_ptr());
10143        let p2 = SendMut(o2.as_mut_ptr());
10144        let run_range = |start: usize, end: usize| {
10145            q8_range2_sdot(q, row_scale, &a1s, &a2s, cols, p1, p2, start, end)
10146        };
10147        match pool {
10148            Some(pool) if rows >= 256 => pool.run_rows(rows, &run_range),
10149            _ => run_range(0, rows),
10150        }
10151        return;
10152    }
10153    #[cfg(target_arch = "x86_64")]
10154    if avx2_a8w8_enabled() {
10155        let a1s = if dtype == TensorDtype::Q8_2f {
10156            split_act_q8_2f(x1, col_field)
10157        } else {
10158            split_act(x1)
10159        };
10160        let a2s = if dtype == TensorDtype::Q8_2f {
10161            split_act_q8_2f(x2, col_field)
10162        } else {
10163            split_act(x2)
10164        };
10165        let p1 = SendMut(o1.as_mut_ptr());
10166        let p2 = SendMut(o2.as_mut_ptr());
10167        let run_range = |start: usize, end: usize| {
10168            q8_range2_avx2(q, row_scale, &a1s, &a2s, cols, p1, p2, start, end)
10169        };
10170        match pool {
10171            Some(pool) if rows >= 256 => pool.run_rows(rows, &run_range),
10172            _ => run_range(0, rows),
10173        }
10174        return;
10175    }
10176
10177    prescale_with(x1, col_field, dtype, 1, |x1s| {
10178        prescale_with(x2, col_field, dtype, 2, |x2s| {
10179            let p1 = SendMut(o1.as_mut_ptr());
10180            let p2 = SendMut(o2.as_mut_ptr());
10181            let run_range = move |start: usize, end: usize| {
10182                for o in start..end {
10183                    let row = &q[o * cols..(o + 1) * cols];
10184                    let s1 = dot_i8_f32(row, x1s) * row_scale[o];
10185                    let s2 = dot_i8_f32(row, x2s) * row_scale[o];
10186                    // SAFETY: disjoint row ranges per worker.
10187                    unsafe {
10188                        *p1.at(o) = s1;
10189                        *p2.at(o) = s2;
10190                    }
10191                }
10192            };
10193            match pool {
10194                Some(pool) if rows >= 256 => pool.run_rows(rows, &run_range),
10195                _ => run_range(0, rows),
10196            }
10197        });
10198    });
10199}
10200
10201#[derive(Clone, Copy)]
10202struct SendMut(*mut f32);
10203unsafe impl Send for SendMut {}
10204unsafe impl Sync for SendMut {}
10205
10206impl SendMut {
10207    #[inline]
10208    fn at(self, i: usize) -> *mut f32 {
10209        unsafe { self.0.add(i) }
10210    }
10211}
10212
10213#[cfg(test)]
10214mod tests {
10215    use super::*;
10216
10217    #[test]
10218    fn q2tp_i8_dot_matches_exact_on_grid() {
10219        // On-grid activations (±1 → sx=1/127, xq=±127 dequantizes
10220        // exactly, no outliers) must make the integer path agree with
10221        // the exact scalar walk to f32 rounding.
10222        let (rows, cols) = (5, 64);
10223        let gpr = cols / GROUP_SIZE;
10224        // Synthetic codes plane + a flat ladder: scales_into is not under
10225        // test here, so drive dot_q2tp_row_i8 / q2tp_row_exact directly
10226        // with hand-made scales.
10227        let chunks: Vec<u8> = (0..rows * gpr * Q2TP_CHUNK)
10228            .map(|i| (i as u32).wrapping_mul(2654435761) as u8)
10229            .collect();
10230        let scales: Vec<f32> = (0..gpr).map(|g| 0.5 + g as f32 * 0.25).collect();
10231        let x: Vec<f32> = (0..cols)
10232            .map(|i| if i % 3 == 0 { -1.0 } else { 1.0 })
10233            .collect();
10234        let act = split_act(&x);
10235        assert!(
10236            act.outliers.is_empty(),
10237            "on-grid input must have no outliers"
10238        );
10239        let gsum = q1_group_sums(&act.xq, gpr);
10240        for r in 0..rows {
10241            let exact = q2tp_row_exact(&chunks, r, gpr, &x, &scales);
10242            let fast = dot_q2tp_row_i8(&chunks, r, gpr, &act.xq, &gsum, &scales) * act.sx;
10243            assert!(
10244                (exact - fast).abs() <= exact.abs() * 1e-5 + 1e-5,
10245                "row {r}: exact {exact} vs i8 {fast}"
10246            );
10247        }
10248    }
10249
10250    #[test]
10251    fn q8_row_dot_fast_matches_scalar() {
10252        // The per-arch fast dot must agree with the exact scalar oracle
10253        // (same contract the fused q8 FFN arm rides on).
10254        let cols = 96;
10255        let row: Vec<u8> = (0..cols)
10256            .map(|i| ((i * 37 % 251) - 125) as i8 as u8)
10257            .collect();
10258        let x: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.13).sin()).collect();
10259        let act = split_act(&x);
10260        let fast = q8_row_dot(&row, &act);
10261        let scalar = q8_row_dot_scalar(&row, &act);
10262        assert!(
10263            (fast - scalar).abs() <= scalar.abs() * 1e-5 + 1e-5,
10264            "fast {fast} vs scalar {scalar}"
10265        );
10266    }
10267
10268    #[test]
10269    fn f32_matvec_matches_matvec_rows_bitexact() {
10270        let (rows, cols) = (300, 40);
10271        let w: Vec<f32> = (0..rows * cols).map(|i| (i as f32 * 0.017).sin()).collect();
10272        let x: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.05).cos()).collect();
10273        let qt = QTensor::from_f32(w.clone(), rows, cols);
10274
10275        let mut a = vec![0.0f32; rows];
10276        matvec_rows(None, &w, &x, &mut a);
10277        let mut b = vec![0.0f32; rows];
10278        qt.matvec(&x, &mut b, None);
10279        assert_eq!(a, b);
10280    }
10281
10282    #[test]
10283    fn sdot_kernel_exact_on_grid() {
10284        // Activations already on the i8 grid (±1 with amax=1 → sx=1/127,
10285        // xq=±127 dequantizes EXACTLY) → the SDOT path must match the
10286        // exact f32 dot to float rounding. This isolates kernel
10287        // correctness from quantization noise.
10288        eprintln!("sdot_enabled = {}", sdot_enabled());
10289        let (rows, cols) = (9, 80); // odd rows → exercises 4-row + tail
10290        let w: Vec<u8> = (0..rows * cols)
10291            .map(|i| (((i * 37) % 251) as i32 - 125) as i8 as u8)
10292            .collect();
10293        let scales: Vec<f32> = (0..rows).map(|o| 0.005 + o as f32 * 0.001).collect();
10294        let x: Vec<f32> = (0..cols)
10295            .map(|i| match i % 3 {
10296                0 => 1.0,
10297                1 => -1.0,
10298                _ => 0.0,
10299            })
10300            .collect();
10301        let mut a = vec![0.0f32; rows];
10302        qmatvec(
10303            &w,
10304            &[],
10305            &scales,
10306            &x,
10307            &[],
10308            TensorDtype::Q8Row,
10309            rows,
10310            cols,
10311            &mut a,
10312            None,
10313        );
10314        for o in 0..rows {
10315            let mut acc = 0.0f32;
10316            for j in 0..cols {
10317                acc += (w[o * cols + j] as i8) as f32 * x[j];
10318            }
10319            let expect = acc * scales[o];
10320            assert!(
10321                (a[o] - expect).abs() < 1e-3 * expect.abs().max(1e-3),
10322                "row {o}: {} vs {expect}",
10323                a[o]
10324            );
10325        }
10326    }
10327
10328    #[test]
10329    fn q1_tbl_fast_path_matches_reference() {
10330        // gpr = 8 exercises the TBL pair-load fast loop, and the LAST
10331        // row's final 4-tile window trips the 4B-overread guard (the
10332        // payload ends exactly at the last tile) — both paths must
10333        // agree with the dequant reference.
10334        let (rows, cols) = (5, 256);
10335        let gpr = cols / GROUP_SIZE;
10336        let mut bytes = Vec::new();
10337        for t in 0..rows * gpr {
10338            let s = 0.007 + (t % 11) as f32 * 0.004;
10339            bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
10340            for j in 0..4 {
10341                bytes.push(((t * 53 + j * 89 + 7) % 249) as u8);
10342            }
10343        }
10344        let x: Vec<f32> = (0..cols)
10345            .map(|i| if (i * 5) % 7 < 3 { 1.0 } else { -1.0 })
10346            .collect();
10347        let mut w = vec![0.0f32; rows * cols];
10348        cortiq_core::quant::dequant_q1(&bytes, &mut w);
10349        let mut got = vec![0.0f32; rows];
10350        q1_matvec(&bytes, &x, rows, cols, &mut got, None);
10351        for o in 0..rows {
10352            let expect: f32 = (0..cols).map(|j| w[o * cols + j] * x[j]).sum();
10353            assert!(
10354                (got[o] - expect).abs() < 1e-3 * expect.abs().max(1e-3),
10355                "row {o}: {} vs {expect}",
10356                got[o]
10357            );
10358        }
10359        // Blocked 1×4 batch (b=5: one quad + remainder) must equal the
10360        // single-matvec path bit-for-bit.
10361        let b = 5usize;
10362        let mut xs_all = Vec::new();
10363        for bi in 0..b {
10364            xs_all.extend(x.iter().map(|v| if bi % 2 == 0 { *v } else { -*v }));
10365        }
10366        let mut mm = vec![0.0f32; b * rows];
10367        q1_matmat(&bytes, &xs_all, b, rows, cols, &mut mm, None);
10368        for bi in 0..b {
10369            let mut single = vec![0.0f32; rows];
10370            q1_matvec(
10371                &bytes,
10372                &xs_all[bi * cols..(bi + 1) * cols],
10373                rows,
10374                cols,
10375                &mut single,
10376                None,
10377            );
10378            assert_eq!(&mm[bi * rows..(bi + 1) * rows], &single[..], "stream {bi}");
10379        }
10380    }
10381
10382    #[test]
10383    fn q1_kernels_match_exact_reference() {
10384        // Synthetic q1 payload: 6-byte tiles [f16 scale][4B bits].
10385        let (rows, cols) = (7, 96);
10386        let gpr = cols / GROUP_SIZE;
10387        let mut bytes = Vec::new();
10388        for t in 0..rows * gpr {
10389            let s = 0.01 + (t % 13) as f32 * 0.003;
10390            bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
10391            for j in 0..4 {
10392                bytes.push(((t * 31 + j * 97) % 251) as u8);
10393            }
10394        }
10395        // On-grid activations (±1, amax 1) → the SDOT path is exact.
10396        let x: Vec<f32> = (0..cols)
10397            .map(|i| if i % 3 == 0 { 1.0 } else { -1.0 })
10398            .collect();
10399        // Reference through the core dequant.
10400        let mut w = vec![0.0f32; rows * cols];
10401        cortiq_core::quant::dequant_q1(&bytes, &mut w);
10402        let mut expect = vec![0.0f32; rows];
10403        for o in 0..rows {
10404            expect[o] = (0..cols).map(|j| w[o * cols + j] * x[j]).sum();
10405        }
10406        let mut got = vec![0.0f32; rows];
10407        q1_matvec(&bytes, &x, rows, cols, &mut got, None);
10408        for o in 0..rows {
10409            assert!(
10410                (got[o] - expect[o]).abs() < 1e-3 * expect[o].abs().max(1e-3),
10411                "row {o}: {} vs {}",
10412                got[o],
10413                expect[o]
10414            );
10415        }
10416        // Pair and batch paths agree with the single path.
10417        let x2: Vec<f32> = x.iter().map(|v| -v).collect();
10418        let (mut a1, mut a2) = (vec![0.0f32; rows], vec![0.0f32; rows]);
10419        q1_matvec2(&bytes, &x, &x2, rows, cols, &mut a1, &mut a2, None);
10420        assert_eq!(a1, got);
10421        let mut xs = x.clone();
10422        xs.extend_from_slice(&x2);
10423        let mut mm = vec![0.0f32; 2 * rows];
10424        q1_matmat(&bytes, &xs, 2, rows, cols, &mut mm, None);
10425        assert_eq!(&mm[..rows], got.as_slice());
10426        assert_eq!(&mm[rows..], a2.as_slice());
10427    }
10428
10429    #[test]
10430    fn repack_is_bit_identical() {
10431        // The interleaved-repack kernel must produce EXACTLY the same
10432        // bits as the mmap-layout kernel: integer accumulation is order-
10433        // exact, the f32 epilogue is identical. Odd rows exercise the
10434        // tail; direct range calls exercise unaligned pool splits.
10435        let (rows, cols) = (267, 96); // 66 groups + 3 tail rows, cols % 16 == 0
10436        let w: Vec<u8> = (0..rows * cols)
10437            .map(|i| (((i * 89) % 253) as i32 - 126) as i8 as u8)
10438            .collect();
10439        let scales: Vec<f32> = (0..rows).map(|o| 0.003 + o as f32 * 0.0007).collect();
10440        let x: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.37).sin() * 2.0).collect();
10441        let rep = q8_repack_layout(&w, rows, cols);
10442        // Group interleave round-trips.
10443        for g in 0..rows / 4 {
10444            for c in 0..cols / 16 {
10445                for lane in 0..4 {
10446                    assert_eq!(
10447                        &rep[g * 4 * cols + c * 64 + lane * 16
10448                            ..g * 4 * cols + c * 64 + lane * 16 + 16],
10449                        &w[(g * 4 + lane) * cols + c * 16..(g * 4 + lane) * cols + c * 16 + 16],
10450                    );
10451                }
10452            }
10453        }
10454        let mut a = vec![0.0f32; rows];
10455        qmatvec(
10456            &w,
10457            &[],
10458            &scales,
10459            &x,
10460            &[],
10461            TensorDtype::Q8Row,
10462            rows,
10463            cols,
10464            &mut a,
10465            None,
10466        );
10467        let mut b = vec![0.0f32; rows];
10468        qmatvec(
10469            &w,
10470            &rep,
10471            &scales,
10472            &x,
10473            &[],
10474            TensorDtype::Q8Row,
10475            rows,
10476            cols,
10477            &mut b,
10478            None,
10479        );
10480        assert_eq!(a, b, "full-range repack output diverged");
10481
10482        #[cfg(target_arch = "aarch64")]
10483        if sdot_enabled() {
10484            // Unaligned range split (pool workers get arbitrary bounds).
10485            let act = split_act(&x);
10486            let mut c1 = vec![0.0f32; rows];
10487            let mut c2 = vec![0.0f32; rows];
10488            q8_range_sdot(
10489                &w,
10490                &[],
10491                &scales,
10492                &act,
10493                cols,
10494                SendMut(c1.as_mut_ptr()),
10495                3,
10496                rows - 2,
10497            );
10498            q8_range_sdot(
10499                &w,
10500                &rep,
10501                &scales,
10502                &act,
10503                cols,
10504                SendMut(c2.as_mut_ptr()),
10505                3,
10506                rows - 2,
10507            );
10508            assert_eq!(c1, c2, "unaligned-range repack output diverged");
10509        }
10510    }
10511
10512    #[test]
10513    fn sdot_a8w8_noise_is_bounded() {
10514        // Off-grid activations: A8 quantization noise must stay small in
10515        // relative L2 over the whole output (realistic accuracy contract;
10516        // vmfcore measured argmax-identical decode on real models).
10517        let (rows, cols) = (16, 512);
10518        let w: Vec<u8> = (0..rows * cols)
10519            .map(|i| (((i * 37) % 251) as i32 - 125) as i8 as u8)
10520            .collect();
10521        let scales = vec![0.01f32; rows];
10522        let x: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.21).sin()).collect();
10523        let mut a = vec![0.0f32; rows];
10524        qmatvec(
10525            &w,
10526            &[],
10527            &scales,
10528            &x,
10529            &[],
10530            TensorDtype::Q8Row,
10531            rows,
10532            cols,
10533            &mut a,
10534            None,
10535        );
10536        let (mut num, mut den) = (0f64, 0f64);
10537        for o in 0..rows {
10538            let mut acc = 0.0f32;
10539            for j in 0..cols {
10540                acc += (w[o * cols + j] as i8) as f32 * x[j];
10541            }
10542            let expect = acc * scales[o];
10543            num += ((a[o] - expect) as f64).powi(2);
10544            den += (expect as f64).powi(2);
10545        }
10546        let rel = (num / den.max(1e-12)).sqrt();
10547        assert!(rel < 0.05, "A8W8 relative L2 error too high: {rel}");
10548    }
10549
10550    #[test]
10551    fn i8_dot_neon_matches_scalar() {
10552        let n = 100;
10553        let w: Vec<u8> = (0..n).map(|i| ((i * 37 + 11) % 251) as u8).collect();
10554        let x: Vec<f32> = (0..n).map(|i| (i as f32 * 0.13).sin()).collect();
10555        let mut scalar = 0.0f32;
10556        for j in 0..n {
10557            scalar += (w[j] as i8) as f32 * x[j];
10558        }
10559        let fast = dot_i8_f32(&w, &x);
10560        assert!((scalar - fast).abs() < 1e-3 * scalar.abs().max(1.0));
10561    }
10562
10563    /// Fused vbit matvec must match full dequant_vbit + dense matvec.
10564    #[test]
10565    fn vbitmatvec_matches_full_dequant() {
10566        let (rows, cols) = (6, 64);
10567        let ng = cols / GROUP_SIZE;
10568        // Hand-craft: bits per row, f16 scales, packed rows.
10569        let bits: Vec<u8> = vec![3, 4, 5, 6, 8, 4];
10570        let mut bytes = bits.clone();
10571        for g in 0..rows * ng {
10572            let s = 0.02 + 0.001 * g as f32;
10573            bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
10574        }
10575        for r in 0..rows {
10576            let b = bits[r] as usize;
10577            let (mut acc, mut nb) = (0u64, 0usize);
10578            let mut rowbytes = Vec::new();
10579            for i in 0..cols {
10580                let v = ((i * 7 + r * 13) % (1 << b)) as u64;
10581                acc = (acc << b) | v;
10582                nb += b;
10583                while nb >= 8 {
10584                    nb -= 8;
10585                    rowbytes.push(((acc >> nb) & 0xFF) as u8);
10586                }
10587            }
10588            if nb > 0 {
10589                rowbytes.push(((acc << (8 - nb)) & 0xFF) as u8);
10590            }
10591            bytes.extend_from_slice(&rowbytes);
10592        }
10593        let x: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.19).sin()).collect();
10594
10595        let mut reference = vec![0f32; rows * cols];
10596        cortiq_core::quant::dequant_vbit(&bytes, rows, cols, &mut reference).unwrap();
10597        let mut expect = vec![0f32; rows];
10598        for r in 0..rows {
10599            expect[r] = reference[r * cols..(r + 1) * cols]
10600                .iter()
10601                .zip(&x)
10602                .map(|(w, xv)| w * xv)
10603                .sum();
10604        }
10605        let mut got = vec![0f32; rows];
10606        let offsets = vbit_row_offsets(&bytes, rows, cols);
10607        vbitmatvec(&bytes, &offsets, &x, rows, cols, &mut got, None);
10608        // SDOT path quantizes activations to i8 (A8W8): bounded noise,
10609        // same contract as q8 (exact path is pinned by CMF_SDOT=0 in
10610        // the golden-parity gate).
10611        let tol = if a8w8_enabled() { 6e-2 } else { 1e-4 };
10612        let scale = expect.iter().fold(0f32, |m, v| m.max(v.abs())).max(1e-6);
10613        for r in 0..rows {
10614            assert!(
10615                (got[r] - expect[r]).abs() < tol * scale,
10616                "row {r}: {} vs {}",
10617                got[r],
10618                expect[r]
10619            );
10620        }
10621    }
10622
10623    /// Fused q4 matvec must match the reference full-dequant + dense
10624    /// matvec bit-for-bit in structure (same f32 math, group order).
10625    /// vbit matmat: the blocked 1×4 leg must match the per-row path
10626    /// (paired env toggle; larger shape so both code paths engage).
10627    #[test]
10628    #[cfg(target_arch = "x86_64")]
10629    fn vbit_matmat_blocked_matches_per_row() {
10630        let (rows, cols, b) = (64usize, 128usize, 9usize);
10631        let ng = cols / GROUP_SIZE;
10632        let bits: Vec<u8> = (0..rows).map(|r| [3u8, 4, 5, 6][r % 4]).collect();
10633        let mut bytes = bits.clone();
10634        for g in 0..rows * ng {
10635            let sc = 0.02 + 0.0005 * g as f32;
10636            bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(sc).to_le_bytes());
10637        }
10638        for r in 0..rows {
10639            let bw = bits[r] as usize;
10640            let (mut acc, mut nb) = (0u64, 0usize);
10641            let mut rowbytes = Vec::new();
10642            for i in 0..cols {
10643                let v = ((i * 7 + r * 13) % (1 << bw)) as u64;
10644                acc = (acc << bw) | v;
10645                nb += bw;
10646                while nb >= 8 {
10647                    nb -= 8;
10648                    rowbytes.push(((acc >> nb) & 0xFF) as u8);
10649                }
10650            }
10651            if nb > 0 {
10652                rowbytes.push(((acc << (8 - nb)) & 0xFF) as u8);
10653            }
10654            bytes.extend_from_slice(&rowbytes);
10655        }
10656        let x: Vec<f32> = (0..b * cols)
10657            .map(|i| ((i * 13 + 7) % 97) as f32 / 97.0 - 0.5)
10658            .collect();
10659        let offsets = vbit_row_offsets(&bytes, rows, cols);
10660        let mut y_a = vec![0f32; b * rows];
10661        let mut y_b = vec![0f32; b * rows];
10662        unsafe { std::env::set_var("CMF_X86_BLOCKED", "1") };
10663        vbitmatmat(&bytes, &offsets, &x, b, rows, cols, &mut y_a, None);
10664        unsafe { std::env::set_var("CMF_X86_BLOCKED", "0") };
10665        vbitmatmat(&bytes, &offsets, &x, b, rows, cols, &mut y_b, None);
10666        unsafe { std::env::remove_var("CMF_X86_BLOCKED") };
10667        let max_d = y_a
10668            .iter()
10669            .zip(&y_b)
10670            .map(|(p, q)| (p - q).abs())
10671            .fold(0.0f32, f32::max);
10672        assert!(max_d < 1e-4, "vbit blocked ≠ per-row: max|Δ| = {max_d}");
10673    }
10674
10675    /// q4t blocked 1×4 (SDOT on ARM, AVX2 on x86) must equal the
10676    /// per-row path exactly: same nibble unpack, same group order,
10677    /// same f32 accumulation — batch == matvec bit-for-bit. b=9 covers
10678    /// two full 1×4 blocks plus a remainder through the single-row
10679    /// kernel. (Both paths produce identical output, so the shared
10680    /// CMF_X86_BLOCKED env var racing with other tests cannot flip
10681    /// the verdict — worst case both sides take the same path.)
10682    #[test]
10683    fn q4t_matmat_blocked_matches_per_row() {
10684        let (rows, cols, b) = (16usize, 64usize, 9usize);
10685        let gpr = cols / GROUP_SIZE;
10686        let mut bytes = vec![0u8; rows * gpr * Q4_TILE];
10687        for r in 0..rows {
10688            for g in 0..gpr {
10689                let t = (r * gpr + g) * Q4_TILE;
10690                let sc = 0.02 + 0.001 * (r * gpr + g) as f32;
10691                bytes[t..t + 2].copy_from_slice(&cortiq_core::quant::f32_to_f16(sc).to_le_bytes());
10692                for k in 0..16 {
10693                    bytes[t + 2 + k] = ((r * 31 + g * 7 + k * 13) % 251) as u8;
10694                }
10695            }
10696        }
10697        let x: Vec<f32> = (0..b * cols)
10698            .map(|i| ((i * 13 + 7) % 97) as f32 / 97.0 - 0.5)
10699            .collect();
10700        let mut y_blk = vec![0f32; b * rows];
10701        let mut y_row = vec![0f32; b * rows];
10702        unsafe { std::env::set_var("CMF_X86_BLOCKED", "1") };
10703        q4t_matmat(&bytes, &x, b, rows, cols, &mut y_blk, None);
10704        unsafe { std::env::set_var("CMF_X86_BLOCKED", "0") };
10705        q4t_matmat(&bytes, &x, b, rows, cols, &mut y_row, None);
10706        unsafe { std::env::remove_var("CMF_X86_BLOCKED") };
10707        assert_eq!(y_blk, y_row, "q4t blocked 1x4 ≠ per-row");
10708    }
10709
10710    /// The wide-batch Accelerate arm of q4t_matmat vs a brute-force
10711    /// f32 dequant matmul: both are f32 GEMMs, so only reduction
10712    /// order differs — tight tolerance.
10713    /// A synthetic q4tp payload: random nibbles plus a per-row ladder whose
10714    /// span varies row to row, so the codes actually exercise the full 0..31
10715    /// range rather than clustering on one rung.
10716    fn synth_q4tp(rows: usize, cols: usize) -> Vec<u8> {
10717        use cortiq_core::quant::{f32_to_f16, q4tp_code_stride, q4tp_put_code};
10718        let gpr = cols / GROUP_SIZE;
10719        let stride = q4tp_code_stride(gpr);
10720        let (params_off, codes_off, _) = q4tp_sections(rows, cols);
10721        let mut b = vec![0u8; codes_off + rows * stride];
10722        for r in 0..rows {
10723            for g in 0..gpr {
10724                let t = (r * gpr + g) * Q4TP_NIB;
10725                for k in 0..16 {
10726                    b[t + k] = ((r * 31 + g * 7 + k * 13) % 251) as u8;
10727                }
10728            }
10729            let lo = -6.0 - 0.03 * (r % 17) as f32;
10730            let step = 0.01 + 0.004 * (r % 11) as f32;
10731            let p = params_off + r * 4;
10732            b[p..p + 2].copy_from_slice(&f32_to_f16(lo).to_le_bytes());
10733            b[p + 2..p + 4].copy_from_slice(&f32_to_f16(step).to_le_bytes());
10734            let crow = &mut b[codes_off + r * stride..codes_off + (r + 1) * stride];
10735            for g in 0..gpr {
10736                q4tp_put_code(crow, g, (r * 5 + g * 3) % 32);
10737            }
10738        }
10739        b
10740    }
10741
10742    /// The same weights re-expressed as q4_tiled, so the proven kernel can
10743    /// be the reference: each tile stores the ladder scale its code selects.
10744    /// Only the f16 rounding of that scale separates the two payloads.
10745    fn q4tp_as_q4t(bytes: &[u8], rows: usize, cols: usize) -> Vec<u8> {
10746        let gpr = cols / GROUP_SIZE;
10747        let v = Q4tpView::new(bytes, rows, cols);
10748        let mut out = vec![0u8; rows * gpr * Q4_TILE];
10749        let mut sc = vec![0f32; gpr];
10750        for r in 0..rows {
10751            v.scales_into(r, gpr, &mut sc);
10752            for g in 0..gpr {
10753                let t = (r * gpr + g) * Q4_TILE;
10754                let s = sc[g];
10755                out[t..t + 2].copy_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
10756                let src = (r * gpr + g) * Q4TP_NIB;
10757                out[t + 2..t + Q4_TILE].copy_from_slice(&v.nib[src..src + Q4TP_NIB]);
10758            }
10759        }
10760        out
10761    }
10762
10763    /// The exact (`CMF_SDOT=0`) path must reproduce `dequant_q4tp` to f32
10764    /// rounding — that scalar routine is the format's definition, and the
10765    /// kernels re-derive the scale from the ladder independently. Call the
10766    /// row kernel directly: `matmat` picks the int8 arm when a8w8 is on,
10767    /// so routing through it would test the other path by accident.
10768    #[test]
10769    fn q4tp_exact_path_matches_dequant_reference() {
10770        let (rows, cols) = (256usize, 512usize);
10771        let gpr = cols / GROUP_SIZE;
10772        let bytes = synth_q4tp(rows, cols);
10773        let mut w = vec![0f32; rows * cols];
10774        cortiq_core::quant::dequant_q4tp(&bytes, rows, cols, &mut w);
10775
10776        let x: Vec<f32> = (0..cols)
10777            .map(|i| ((i * 13 + 7) % 97) as f32 / 97.0 - 0.5)
10778            .collect();
10779        let v = Q4tpView::new(&bytes, rows, cols);
10780        let mut sc = vec![0f32; gpr];
10781        for r in 0..rows {
10782            v.scales_into(r, gpr, &mut sc);
10783            let got = q4tp_row_exact(v.nib, r, gpr, &x, &sc);
10784            let want: f32 = (0..cols).map(|c| w[r * cols + c] * x[c]).sum();
10785            // These dot products cancel down to ~1e-3 from terms of ~5e-2, so
10786            // the meaningful yardstick is the summed magnitude, not the result:
10787            // against the result any reordering of a 512-term f32 sum "fails".
10788            let mag: f32 = (0..cols).map(|c| (w[r * cols + c] * x[c]).abs()).sum();
10789            assert!(
10790                (got - want).abs() <= 1e-5 * mag,
10791                "row {r}: kernel {got} vs dequant {want}"
10792            );
10793        }
10794    }
10795
10796    /// The int8 (a8w8) path can't be checked against an f32 reference — the
10797    /// activation quantization dominates. Check it against the q4t kernel it
10798    /// was ported from instead, on payloads holding the same weights: that
10799    /// isolates exactly what the port could break (16 B stride, ladder
10800    /// lookup, nibble unpack) from what it deliberately shares.
10801    #[test]
10802    fn q4tp_matvec_matches_the_q4t_kernel_it_was_ported_from() {
10803        let (rows, cols) = (256usize, 512usize);
10804        let bytes = synth_q4tp(rows, cols);
10805        let twin = q4tp_as_q4t(&bytes, rows, cols);
10806        let x: Vec<f32> = (0..cols)
10807            .map(|i| ((i * 13 + 7) % 97) as f32 / 97.0 - 0.5)
10808            .collect();
10809
10810        let mut got = vec![0f32; rows];
10811        q4tp_matvec(&bytes, &x, rows, cols, &mut got, None);
10812        let mut want = vec![0f32; rows];
10813        q4t_matvec(&twin, &x, rows, cols, &mut want, None);
10814
10815        // Scale is f16 in the twin and f32 here, so allow that rounding on
10816        // top of the summed magnitude (same cancellation argument as above).
10817        let mut w = vec![0f32; rows * cols];
10818        cortiq_core::quant::dequant_q4tp(&bytes, rows, cols, &mut w);
10819        for r in 0..rows {
10820            let mag: f32 = (0..cols).map(|c| (w[r * cols + c] * x[c]).abs()).sum();
10821            assert!(
10822                (got[r] - want[r]).abs() <= 1e-3 * mag,
10823                "row {r}: q4tp {} vs q4t {}",
10824                got[r],
10825                want[r]
10826            );
10827        }
10828    }
10829
10830    /// `matmat` carries three arms (Accelerate, blocked int8 1x4, scalar).
10831    /// Batch 5 crosses the blocked kernel's stride, so this exercises the
10832    /// 1x4 path AND its scalar tail in one run — the blocked kernel is new
10833    /// code and its four accumulators are exactly what tends to go wrong.
10834    #[test]
10835    fn q4tp_matmat_matches_the_q4t_kernel_it_was_ported_from() {
10836        let (rows, cols, b) = (256usize, 512usize, 5usize);
10837        let bytes = synth_q4tp(rows, cols);
10838        let twin = q4tp_as_q4t(&bytes, rows, cols);
10839        let xs: Vec<f32> = (0..b * cols)
10840            .map(|i| ((i * 29 + 11) % 89) as f32 / 89.0 - 0.5)
10841            .collect();
10842
10843        let mut got = vec![0f32; b * rows];
10844        q4tp_matmat(&bytes, &xs, b, rows, cols, &mut got, None);
10845        let mut want = vec![0f32; b * rows];
10846        q4t_matmat(&twin, &xs, b, rows, cols, &mut want, None);
10847
10848        let mut w = vec![0f32; rows * cols];
10849        cortiq_core::quant::dequant_q4tp(&bytes, rows, cols, &mut w);
10850        for t in 0..b {
10851            for r in 0..rows {
10852                let mag: f32 = (0..cols)
10853                    .map(|c| (w[r * cols + c] * xs[t * cols + c]).abs())
10854                    .sum();
10855                let (g, wa) = (got[t * rows + r], want[t * rows + r]);
10856                assert!(
10857                    (g - wa).abs() <= 1e-3 * mag,
10858                    "batch {t} row {r}: q4tp {g} vs q4t {wa}"
10859                );
10860            }
10861        }
10862    }
10863
10864    #[test]
10865    fn q4tp_matvec2_matches_the_single_stream_kernel() {
10866        let (rows, cols) = (128usize, 256usize);
10867        let gpr = cols / GROUP_SIZE;
10868        let bytes = synth_q4tp(rows, cols);
10869        let xs: Vec<f32> = (0..2 * cols)
10870            .map(|i| ((i * 29 + 11) % 89) as f32 / 89.0 - 0.5)
10871            .collect();
10872
10873        let (mut o1, mut o2) = (vec![0f32; rows], vec![0f32; rows]);
10874        q4tp_matvec2(
10875            &bytes,
10876            &xs[..cols],
10877            &xs[cols..],
10878            rows,
10879            cols,
10880            &mut o1,
10881            &mut o2,
10882            None,
10883        );
10884
10885        // matvec2 takes the exact path for both streams, so the single-row
10886        // kernel is an exact reference — no tolerance for path differences.
10887        let v = Q4tpView::new(&bytes, rows, cols);
10888        let mut sc = vec![0f32; gpr];
10889        for r in 0..rows {
10890            v.scales_into(r, gpr, &mut sc);
10891            assert_eq!(o1[r], q4tp_row_exact(v.nib, r, gpr, &xs[..cols], &sc));
10892            assert_eq!(o2[r], q4tp_row_exact(v.nib, r, gpr, &xs[cols..], &sc));
10893        }
10894    }
10895
10896    /// q4tp must not COST speed — it exists to save bytes, and a format that
10897    /// trades 7% of a file for a slower model is a bad trade. This guard is
10898    /// here because correctness tests happily passed while `q4tp_matmat` was
10899    /// missing its int8 and Accelerate arms and the model ran 5x slower.
10900    /// Measured on M-series: 0.97-1.04x, i.e. parity (16 B tiles are better
10901    /// aligned than q4t's 18 B, which pays for the scale indirection).
10902    #[test]
10903    fn q4tp_matvec_keeps_pace_with_q4t() {
10904        let (rows, cols) = (4096usize, 3072usize);
10905        let bytes = synth_q4tp(rows, cols);
10906        let twin = q4tp_as_q4t(&bytes, rows, cols);
10907        let x: Vec<f32> = (0..cols).map(|i| (i % 97) as f32 / 97.0 - 0.5).collect();
10908        let mut o = vec![0f32; rows];
10909        let n = 12;
10910        let mut best = (f64::MAX, f64::MAX);
10911        // Interleaved A/B, minimum statistic: this machine throttles, and a
10912        // mean over a thermal ramp reliably indicts whichever ran second.
10913        for _ in 0..3 {
10914            let t0 = std::time::Instant::now();
10915            for _ in 0..n {
10916                q4t_matvec(&twin, &x, rows, cols, &mut o, None);
10917            }
10918            best.0 = best.0.min(t0.elapsed().as_secs_f64());
10919            let t0 = std::time::Instant::now();
10920            for _ in 0..n {
10921                q4tp_matvec(&bytes, &x, rows, cols, &mut o, None);
10922            }
10923            best.1 = best.1.min(t0.elapsed().as_secs_f64());
10924        }
10925        let ratio = best.1 / best.0;
10926        println!(
10927            "q4t {:.3} ms | q4tp {:.3} ms | {ratio:.2}x",
10928            best.0 * 1e3 / n as f64,
10929            best.1 * 1e3 / n as f64
10930        );
10931        assert!(ratio < 2.0, "q4tp matvec {ratio:.2}x slower than q4t");
10932    }
10933
10934    #[cfg(target_os = "macos")]
10935    #[test]
10936    fn q4t_matmat_accel_matches_dequant_reference() {
10937        if !accel_gemm_enabled() {
10938            return; // CMF_ACCEL=0
10939        }
10940        let (rows, cols, b) = (512usize, 1024usize, 8usize); // ≥500K → accel arm
10941        let gpr = cols / GROUP_SIZE;
10942        let mut bytes = vec![0u8; rows * gpr * Q4_TILE];
10943        for r in 0..rows {
10944            for g in 0..gpr {
10945                let t = (r * gpr + g) * Q4_TILE;
10946                let sc = 0.02 + 0.0005 * ((r * gpr + g) % 64) as f32;
10947                bytes[t..t + 2].copy_from_slice(&cortiq_core::quant::f32_to_f16(sc).to_le_bytes());
10948                for k in 0..16 {
10949                    bytes[t + 2 + k] = ((r * 31 + g * 7 + k * 13) % 251) as u8;
10950                }
10951            }
10952        }
10953        let x: Vec<f32> = (0..b * cols)
10954            .map(|i| ((i * 13 + 7) % 97) as f32 / 97.0 - 0.5)
10955            .collect();
10956        let mut got = vec![0f32; b * rows];
10957        q4t_matmat(&bytes, &x, b, rows, cols, &mut got, None);
10958        // Brute-force reference off the same tiles.
10959        let mut w = vec![0f32; rows * cols];
10960        for r in 0..rows {
10961            for g in 0..gpr {
10962                let t = (r * gpr + g) * Q4_TILE;
10963                let s = f16_to_f32(u16::from_le_bytes([bytes[t], bytes[t + 1]]));
10964                for (k, &bb) in bytes[t + 2..t + Q4_TILE].iter().enumerate() {
10965                    w[r * cols + g * GROUP_SIZE + k * 2] = ((bb & 0x0F) as f32 - 8.0) * s;
10966                    w[r * cols + g * GROUP_SIZE + k * 2 + 1] =
10967                        (((bb >> 4) & 0x0F) as f32 - 8.0) * s;
10968                }
10969            }
10970        }
10971        for bi in 0..b {
10972            for r in 0..rows {
10973                let want: f32 = (0..cols).map(|j| x[bi * cols + j] * w[r * cols + j]).sum();
10974                let d = (got[bi * rows + r] - want).abs();
10975                assert!(
10976                    d <= want.abs().max(1.0) * 1e-4,
10977                    "accel q4t GEMM diverged at ({bi},{r}): {} vs {want}",
10978                    got[bi * rows + r]
10979                );
10980            }
10981        }
10982    }
10983
10984    #[test]
10985    fn q4matvec_matches_full_dequant() {
10986        let (rows, cols) = (8, 64);
10987        let groups = rows * cols / GROUP_SIZE;
10988        // Hand-craft a q4_block blob: nibbles then f16 scales.
10989        let mut bytes = Vec::with_capacity(groups * 16 + groups * 2);
10990        for i in 0..groups * 16 {
10991            bytes.push((((i * 7 + 3) % 256) & 0xFF) as u8);
10992        }
10993        for g in 0..groups {
10994            let s = 0.01 + 0.003 * g as f32;
10995            bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
10996        }
10997        let x: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.17).sin()).collect();
10998
10999        let mut reference = vec![0.0f32; rows * cols];
11000        cortiq_core::quant::dequant_q4_block(&bytes, &mut reference);
11001        let mut expect = vec![0.0f32; rows];
11002        for r in 0..rows {
11003            expect[r] = reference[r * cols..(r + 1) * cols]
11004                .iter()
11005                .zip(&x)
11006                .map(|(w, xv)| w * xv)
11007                .sum();
11008        }
11009
11010        let mut got = vec![0.0f32; rows];
11011        q4matvec(&bytes, &x, rows, cols, &mut got, None);
11012        // SDOT path quantizes activations to i8 (A8W8): bounded noise,
11013        // same contract as q8/vbit (exact path is pinned by CMF_SDOT=0
11014        // in the golden-parity gate).
11015        let tol = if a8w8_enabled() { 6e-2 } else { 1e-4 };
11016        let scale = expect.iter().fold(0f32, |m, v| m.max(v.abs())).max(1.0);
11017        for r in 0..rows {
11018            assert!(
11019                (got[r] - expect[r]).abs() < tol * scale,
11020                "row {r}: {} vs {}",
11021                got[r],
11022                expect[r]
11023            );
11024        }
11025    }
11026
11027    /// Fused two-input vbit matvec must equal two single matvecs exactly
11028    /// (same per-lane accumulation order on both scalar and SDOT paths).
11029    #[test]
11030    fn vbitmatvec2_equals_two_singles() {
11031        let (rows, cols) = (6, 64);
11032        let ng = cols / GROUP_SIZE;
11033        let bits: Vec<u8> = vec![3, 4, 5, 6, 8, 4];
11034        let mut bytes = bits.clone();
11035        for g in 0..rows * ng {
11036            let s = 0.02 + 0.001 * g as f32;
11037            bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
11038        }
11039        for r in 0..rows {
11040            let b = bits[r] as usize;
11041            let (mut acc, mut nb) = (0u64, 0usize);
11042            let mut rowbytes = Vec::new();
11043            for i in 0..cols {
11044                let v = ((i * 7 + r * 13) % (1 << b)) as u64;
11045                acc = (acc << b) | v;
11046                nb += b;
11047                while nb >= 8 {
11048                    nb -= 8;
11049                    rowbytes.push(((acc >> nb) & 0xFF) as u8);
11050                }
11051            }
11052            if nb > 0 {
11053                rowbytes.push(((acc << (8 - nb)) & 0xFF) as u8);
11054            }
11055            bytes.extend_from_slice(&rowbytes);
11056        }
11057        let x1: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.19).sin()).collect();
11058        let x2: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.11).cos()).collect();
11059        let offsets = vbit_row_offsets(&bytes, rows, cols);
11060
11061        let (mut a1, mut a2) = (vec![0f32; rows], vec![0f32; rows]);
11062        vbitmatvec(&bytes, &offsets, &x1, rows, cols, &mut a1, None);
11063        vbitmatvec(&bytes, &offsets, &x2, rows, cols, &mut a2, None);
11064        let (mut b1, mut b2) = (vec![0f32; rows], vec![0f32; rows]);
11065        vbitmatvec2(
11066            &bytes, &offsets, &x1, &x2, rows, cols, &mut b1, &mut b2, None,
11067        );
11068        assert_eq!(a1, b1, "fused vbit lane 1 must be bit-identical");
11069        assert_eq!(a2, b2, "fused vbit lane 2 must be bit-identical");
11070    }
11071
11072    /// Fused two-input q4 matvec must equal two single matvecs exactly.
11073    #[test]
11074    fn q4matvec2_equals_two_singles() {
11075        let (rows, cols) = (8, 128);
11076        let groups = rows * cols / GROUP_SIZE;
11077        let mut bytes = Vec::with_capacity(groups * 16 + groups * 2);
11078        for i in 0..groups * 16 {
11079            bytes.push((((i * 7 + 3) % 256) & 0xFF) as u8);
11080        }
11081        for g in 0..groups {
11082            let s = 0.01 + 0.003 * g as f32;
11083            bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
11084        }
11085        // Include an outlier channel so the SDOT correction path is
11086        // exercised in the pair kernel too.
11087        let mut x1: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.17).sin()).collect();
11088        x1[9] = 250.0;
11089        let x2: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.23).cos()).collect();
11090
11091        let (mut a1, mut a2) = (vec![0f32; rows], vec![0f32; rows]);
11092        q4matvec(&bytes, &x1, rows, cols, &mut a1, None);
11093        q4matvec(&bytes, &x2, rows, cols, &mut a2, None);
11094        let (mut b1, mut b2) = (vec![0f32; rows], vec![0f32; rows]);
11095        q4matvec2(&bytes, &x1, &x2, rows, cols, &mut b1, &mut b2, None);
11096        assert_eq!(a1, b1, "fused q4 lane 1 must be bit-identical");
11097        assert_eq!(a2, b2, "fused q4 lane 2 must be bit-identical");
11098    }
11099
11100    /// Multi-matrix job must equal separate matvecs exactly — same
11101    /// kernels, only the dispatch is fused.
11102    #[test]
11103    fn matvec_many_equals_separate_matvecs() {
11104        use crate::pool::Pool;
11105        let (r1, r2, cols) = (300, 200, 64);
11106        let mk = |salt: usize, rows: usize| {
11107            QTensor::from_f32(
11108                (0..rows * cols)
11109                    .map(|i| ((i * 7 + salt) % 97) as f32 / 97.0 - 0.5)
11110                    .collect(),
11111                rows,
11112                cols,
11113            )
11114        };
11115        let (a, b) = (mk(1, r1), mk(5, r2));
11116        let x: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.11).sin()).collect();
11117        let pool = Pool::new(3);
11118
11119        let (mut ea, mut eb) = (vec![0f32; r1], vec![0f32; r2]);
11120        a.matvec(&x, &mut ea, Some(&pool));
11121        b.matvec(&x, &mut eb, Some(&pool));
11122        let (mut ga, mut gb) = (vec![0f32; r1], vec![0f32; r2]);
11123        QTensor::matvec_many([&a, &b], &x, [&mut ga, &mut gb], Some(&pool));
11124        assert_eq!(ea, ga, "fused multi-matrix lane 1 must be bit-identical");
11125        assert_eq!(eb, gb, "fused multi-matrix lane 2 must be bit-identical");
11126    }
11127
11128    /// Batched q4/vbit matmat must equal per-position matvec calls
11129    /// exactly (the fallback it replaced) — same kernels, same order.
11130    #[test]
11131    fn batched_matmat_equals_per_position_matvec() {
11132        let (rows, cols, b) = (8, 64, 5);
11133        // q4 blob.
11134        let groups = rows * cols / GROUP_SIZE;
11135        let mut q4 = Vec::new();
11136        for i in 0..groups * 16 {
11137            q4.push((((i * 7 + 3) % 256) & 0xFF) as u8);
11138        }
11139        for g in 0..groups {
11140            q4.extend_from_slice(
11141                &cortiq_core::quant::f32_to_f16(0.01 + 0.003 * g as f32).to_le_bytes(),
11142            );
11143        }
11144        // vbit blob (mixed widths incl. 8).
11145        let ng = cols / GROUP_SIZE;
11146        let bits: Vec<u8> = vec![3, 4, 5, 6, 8, 4, 5, 3];
11147        let mut vb = bits.clone();
11148        for g in 0..rows * ng {
11149            vb.extend_from_slice(
11150                &cortiq_core::quant::f32_to_f16(0.02 + 0.001 * g as f32).to_le_bytes(),
11151            );
11152        }
11153        for r in 0..rows {
11154            let bw = bits[r] as usize;
11155            let (mut acc, mut nb) = (0u64, 0usize);
11156            let mut rowbytes = Vec::new();
11157            for i in 0..cols {
11158                let v = ((i * 7 + r * 13) % (1 << bw)) as u64;
11159                acc = (acc << bw) | v;
11160                nb += bw;
11161                while nb >= 8 {
11162                    nb -= 8;
11163                    rowbytes.push(((acc >> nb) & 0xFF) as u8);
11164                }
11165            }
11166            if nb > 0 {
11167                rowbytes.push(((acc << (8 - nb)) & 0xFF) as u8);
11168            }
11169            vb.extend_from_slice(&rowbytes);
11170        }
11171        let offsets = vbit_row_offsets(&vb, rows, cols);
11172
11173        let xs: Vec<f32> = (0..b * cols).map(|i| (i as f32 * 0.13).sin()).collect();
11174
11175        // q4: batch vs singles.
11176        let mut got = vec![0f32; b * rows];
11177        q4matmat(&q4, &xs, b, rows, cols, &mut got, None);
11178        for bi in 0..b {
11179            let mut expect = vec![0f32; rows];
11180            q4matvec(
11181                &q4,
11182                &xs[bi * cols..(bi + 1) * cols],
11183                rows,
11184                cols,
11185                &mut expect,
11186                None,
11187            );
11188            assert_eq!(
11189                &got[bi * rows..(bi + 1) * rows],
11190                &expect[..],
11191                "q4 batch pos {bi}"
11192            );
11193        }
11194
11195        // vbit: batch vs singles.
11196        let mut got = vec![0f32; b * rows];
11197        vbitmatmat(&vb, &offsets, &xs, b, rows, cols, &mut got, None);
11198        for bi in 0..b {
11199            let mut expect = vec![0f32; rows];
11200            vbitmatvec(
11201                &vb,
11202                &offsets,
11203                &xs[bi * cols..(bi + 1) * cols],
11204                rows,
11205                cols,
11206                &mut expect,
11207                None,
11208            );
11209            assert_eq!(
11210                &got[bi * rows..(bi + 1) * rows],
11211                &expect[..],
11212                "vbit batch pos {bi}"
11213            );
11214        }
11215    }
11216
11217    /// q4_tiled kernels must produce BIT-identical outputs to the q4
11218    /// split kernels on the same values (same ints, same order — only
11219    /// the byte placement differs).
11220    #[test]
11221    fn q4_tiled_matches_q4_block_bitexact() {
11222        let (rows, cols, b) = (8usize, 128usize, 3usize);
11223        let groups = rows * cols / GROUP_SIZE;
11224        let mut split = Vec::with_capacity(groups * 18);
11225        for i in 0..groups * 16 {
11226            split.push((((i * 7 + 3) % 256) & 0xFF) as u8);
11227        }
11228        for g in 0..groups {
11229            split.extend_from_slice(
11230                &cortiq_core::quant::f32_to_f16(0.01 + 0.003 * g as f32).to_le_bytes(),
11231            );
11232        }
11233        // Re-tile: [scale][nibbles] per group.
11234        let (packed, scales) = split.split_at(groups * 16);
11235        let mut tiled = Vec::with_capacity(groups * Q4_TILE);
11236        for g in 0..groups {
11237            tiled.extend_from_slice(&scales[g * 2..g * 2 + 2]);
11238            tiled.extend_from_slice(&packed[g * 16..(g + 1) * 16]);
11239        }
11240
11241        let mut x1: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.17).sin()).collect();
11242        x1[9] = 250.0; // exercise the outlier path
11243        let x2: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.23).cos()).collect();
11244
11245        let (mut a, mut t) = (vec![0f32; rows], vec![0f32; rows]);
11246        q4matvec(&split, &x1, rows, cols, &mut a, None);
11247        q4t_matvec(&tiled, &x1, rows, cols, &mut t, None);
11248        assert_eq!(a, t, "q4t matvec must match q4 bit-for-bit");
11249
11250        let (mut a1, mut a2) = (vec![0f32; rows], vec![0f32; rows]);
11251        let (mut t1, mut t2) = (vec![0f32; rows], vec![0f32; rows]);
11252        q4matvec2(&split, &x1, &x2, rows, cols, &mut a1, &mut a2, None);
11253        q4t_matvec2(&tiled, &x1, &x2, rows, cols, &mut t1, &mut t2, None);
11254        assert_eq!(a1, t1);
11255        assert_eq!(a2, t2);
11256
11257        let xs: Vec<f32> = (0..b * cols).map(|i| (i as f32 * 0.13).sin()).collect();
11258        let (mut am, mut tm) = (vec![0f32; b * rows], vec![0f32; b * rows]);
11259        q4matmat(&split, &xs, b, rows, cols, &mut am, None);
11260        q4t_matmat(&tiled, &xs, b, rows, cols, &mut tm, None);
11261        assert_eq!(am, tm, "q4t matmat must match q4 bit-for-bit");
11262    }
11263
11264    /// q4 SDOT outlier correction: a single huge activation channel
11265    /// (>8·rms → outlier, zeroed in xq) must still contribute its EXACT
11266    /// term. On-grid bulk (±1/0 → xq dequantizes exactly) isolates the
11267    /// correction from A8W8 noise. cols must exceed 64: at n=64 the
11268    /// 8·rms threshold equals sqrt(v²+rest) ≥ v, so a single outlier
11269    /// can never qualify (8² = n).
11270    #[test]
11271    fn q4matvec_sdot_outlier_exact() {
11272        let (rows, cols) = (4, 128);
11273        let groups = rows * cols / GROUP_SIZE;
11274        let mut bytes = Vec::with_capacity(groups * 16 + groups * 2);
11275        for i in 0..groups * 16 {
11276            bytes.push(((i * 11 + 5) % 256) as u8);
11277        }
11278        for g in 0..groups {
11279            let s = 0.02 + 0.002 * g as f32;
11280            bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
11281        }
11282        let mut x: Vec<f32> = (0..cols)
11283            .map(|i| match i % 3 {
11284                0 => 1.0,
11285                1 => -1.0,
11286                _ => 0.0,
11287            })
11288            .collect();
11289        x[17] = 300.0; // ≫ 8·rms → outlier channel
11290
11291        let mut reference = vec![0.0f32; rows * cols];
11292        cortiq_core::quant::dequant_q4_block(&bytes, &mut reference);
11293        let mut expect = vec![0.0f32; rows];
11294        for r in 0..rows {
11295            expect[r] = reference[r * cols..(r + 1) * cols]
11296                .iter()
11297                .zip(&x)
11298                .map(|(w, xv)| w * xv)
11299                .sum();
11300        }
11301        let mut got = vec![0.0f32; rows];
11302        q4matvec(&bytes, &x, rows, cols, &mut got, None);
11303        let scale = expect.iter().fold(0f32, |m, v| m.max(v.abs())).max(1.0);
11304        for r in 0..rows {
11305            assert!(
11306                (got[r] - expect[r]).abs() < 2e-3 * scale,
11307                "row {r}: {} vs {} (outlier term must be exact)",
11308                got[r],
11309                expect[r]
11310            );
11311        }
11312    }
11313
11314    /// The fused q1t matvec must equal the reference (dequant_q1t → dot),
11315    /// including the ternary zero level and the binary-searched outlier
11316    /// overlay. Guards the mmap kernel that makes a 12B q1t runnable.
11317    #[test]
11318    fn q1t_matvec_matches_reference() {
11319        use cortiq_core::quant::{dequant_q1t, f32_to_f16};
11320        let (rows, cols) = (3usize, 64usize); // gpr = 2
11321        let gpr = cols / GROUP_SIZE;
11322        let scales = [0.5f32, 0.3, 0.7, 0.2, 0.6, 0.15];
11323        // Overlay (must be sorted by flat index): a few spikes across rows.
11324        let outliers: [(u32, f32); 3] = [(5, 9.0), (70, -4.5), (150, 3.25)];
11325        let is_out = |flat: usize| outliers.iter().any(|&(i, _)| i as usize == flat);
11326        let mut bytes = Vec::new();
11327        for r in 0..rows {
11328            for g in 0..gpr {
11329                bytes.extend_from_slice(&f32_to_f16(scales[r * gpr + g]).to_le_bytes());
11330                let mut c = [0u8; 7];
11331                for k in 0..GROUP_SIZE {
11332                    // Encoder invariant: code 0 at outlier positions.
11333                    let code = if is_out(r * cols + g * GROUP_SIZE + k) {
11334                        0
11335                    } else {
11336                        ((k + r * 3 + g) % 3) as u8 // 0,1,2
11337                    };
11338                    cortiq_core::quant::q1t_pack(&mut c, k, code);
11339                }
11340                bytes.extend_from_slice(&c);
11341            }
11342        }
11343        // Per-row overlay: [u32 row_ptr[rows+1]] then [(u16 col, f16 val)] by
11344        // row (outliers are sorted by flat index → already grouped by row).
11345        let mut row_ptr = vec![0u32; rows + 1];
11346        for &(idx, _) in &outliers {
11347            row_ptr[idx as usize / cols + 1] += 1;
11348        }
11349        for r in 0..rows {
11350            row_ptr[r + 1] += row_ptr[r];
11351        }
11352        for &p in &row_ptr {
11353            bytes.extend_from_slice(&p.to_le_bytes());
11354        }
11355        for &(idx, v) in &outliers {
11356            bytes.extend_from_slice(&((idx as usize % cols) as u16).to_le_bytes());
11357            bytes.extend_from_slice(&f32_to_f16(v).to_le_bytes());
11358        }
11359
11360        let mut refw = vec![0f32; rows * cols];
11361        dequant_q1t(&bytes, rows, cols, &mut refw);
11362        // On-grid activations (±1, amax 1) so the int8 SDOT path reconstructs
11363        // x exactly and matches the f32 reference (same trick as the q1 test).
11364        let x: Vec<f32> = (0..cols)
11365            .map(|j| if j % 3 == 0 { 1.0 } else { -1.0 })
11366            .collect();
11367        let mut expect = vec![0f32; rows];
11368        for r in 0..rows {
11369            let mut a = 0.0f32;
11370            for j in 0..cols {
11371                a += refw[r * cols + j] * x[j];
11372            }
11373            expect[r] = a;
11374        }
11375        let tol = |e: f32| 1e-3 * e.abs().max(1e-3);
11376        let mut got = vec![0f32; rows];
11377        q1t_matvec(&bytes, &x, rows, cols, &mut got, None);
11378        for r in 0..rows {
11379            assert!(
11380                (got[r] - expect[r]).abs() < tol(expect[r]),
11381                "row {r}: {} vs {}",
11382                got[r],
11383                expect[r]
11384            );
11385        }
11386        // matmat (b=2, f32 decode path) must agree too.
11387        let x2: Vec<f32> = x.iter().chain(x.iter()).copied().collect();
11388        let mut gm = vec![0f32; 2 * rows];
11389        q1t_matmat(&bytes, &x2, 2, rows, cols, &mut gm, None);
11390        for r in 0..rows {
11391            assert!((gm[r] - expect[r]).abs() < tol(expect[r]));
11392            assert!((gm[rows + r] - expect[r]).abs() < tol(expect[r]));
11393        }
11394        // Fused pair (q1t_matvec2) must equal two single matvecs
11395        // bit-for-bit: same unpack, same group order, same f32
11396        // accumulation per stream. Distinct x2 exercises both lanes.
11397        let xb: Vec<f32> = (0..cols)
11398            .map(|j| if j % 5 == 0 { -1.0 } else { 1.0 })
11399            .collect();
11400        let (mut s1, mut s2) = (vec![0f32; rows], vec![0f32; rows]);
11401        q1t_matvec(&bytes, &x, rows, cols, &mut s1, None);
11402        q1t_matvec(&bytes, &xb, rows, cols, &mut s2, None);
11403        let (mut p1, mut p2) = (vec![0f32; rows], vec![0f32; rows]);
11404        q1t_matvec2(&bytes, &x, &xb, rows, cols, &mut p1, &mut p2, None);
11405        assert_eq!(p1, s1, "q1t pair lane 1 ≠ single matvec");
11406        assert_eq!(p2, s2, "q1t pair lane 2 ≠ single matvec");
11407    }
11408
11409    /// Pair == 2×matvec with an ODD group count (the kernel's tail
11410    /// group) and no overlay section.
11411    #[test]
11412    fn q1t_matvec2_odd_gpr_matches_singles() {
11413        use cortiq_core::quant::{Q1T_TILE, f32_to_f16, q1t_pack};
11414        let (rows, cols) = (5usize, 96usize); // gpr = 3 → paired + tail
11415        let gpr = cols / GROUP_SIZE;
11416        let mut bytes = Vec::with_capacity(rows * gpr * Q1T_TILE);
11417        for r in 0..rows {
11418            for g in 0..gpr {
11419                bytes.extend_from_slice(&f32_to_f16(0.1 + 0.05 * (r + g) as f32).to_le_bytes());
11420                let mut c = [0u8; 7];
11421                for k in 0..GROUP_SIZE {
11422                    q1t_pack(&mut c, k, ((k * 7 + r * 5 + g * 3) % 3) as u8);
11423                }
11424                bytes.extend_from_slice(&c);
11425            }
11426        }
11427        let x1: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.31).sin()).collect();
11428        let x2: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.17).cos()).collect();
11429        let (mut s1, mut s2) = (vec![0f32; rows], vec![0f32; rows]);
11430        q1t_matvec(&bytes, &x1, rows, cols, &mut s1, None);
11431        q1t_matvec(&bytes, &x2, rows, cols, &mut s2, None);
11432        let (mut p1, mut p2) = (vec![0f32; rows], vec![0f32; rows]);
11433        q1t_matvec2(&bytes, &x1, &x2, rows, cols, &mut p1, &mut p2, None);
11434        assert_eq!(p1, s1, "odd-gpr pair lane 1 ≠ single");
11435        assert_eq!(p2, s2, "odd-gpr pair lane 2 ≠ single");
11436    }
11437
11438    // Speed A/B: fused pair (one unpack, two streams) vs two single
11439    // matvecs. Single-threaded, FFN-sized, min-of paired in-process.
11440    //   cargo test -p cortiq-engine --release q1t_matvec2_speed -- --ignored --nocapture
11441    #[test]
11442    #[ignore]
11443    fn q1t_matvec2_speed() {
11444        use cortiq_core::quant::{Q1T_TILE, f32_to_f16, q1t_pack};
11445        use std::time::Instant;
11446        let (rows, cols) = (8192usize, 4096usize);
11447        let gpr = cols / GROUP_SIZE;
11448        let mut bytes = Vec::with_capacity(rows * gpr * Q1T_TILE);
11449        for r in 0..rows {
11450            for g in 0..gpr {
11451                let s = 0.1 + ((r + g) % 7) as f32 * 0.01;
11452                bytes.extend_from_slice(&f32_to_f16(s).to_le_bytes());
11453                let mut c = [0u8; 7];
11454                for k in 0..GROUP_SIZE {
11455                    q1t_pack(&mut c, k, ((k * 7 + r + g) % 3) as u8);
11456                }
11457                bytes.extend_from_slice(&c);
11458            }
11459        }
11460        let x1: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.31).sin()).collect();
11461        let x2: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.17).cos()).collect();
11462        let (mut s1, mut s2) = (vec![0f32; rows], vec![0f32; rows]);
11463        let (mut p1, mut p2) = (vec![0f32; rows], vec![0f32; rows]);
11464        // Warm both paths once.
11465        q1t_matvec(&bytes, &x1, rows, cols, &mut s1, None);
11466        q1t_matvec2(&bytes, &x1, &x2, rows, cols, &mut p1, &mut p2, None);
11467        let (mut t_pair, mut t_two) = (f64::MAX, f64::MAX);
11468        for _ in 0..8 {
11469            let t0 = Instant::now();
11470            q1t_matvec2(&bytes, &x1, &x2, rows, cols, &mut p1, &mut p2, None);
11471            t_pair = t_pair.min(t0.elapsed().as_secs_f64() * 1000.0);
11472            let t1 = Instant::now();
11473            q1t_matvec(&bytes, &x1, rows, cols, &mut s1, None);
11474            q1t_matvec(&bytes, &x2, rows, cols, &mut s2, None);
11475            t_two = t_two.min(t1.elapsed().as_secs_f64() * 1000.0);
11476        }
11477        assert_eq!(p1, s1);
11478        assert_eq!(p2, s2);
11479        println!("q1t pair {rows}x{cols}: fused {t_pair:.2} ms | two singles {t_two:.2} ms");
11480    }
11481
11482    // Speed A/B: the base-3-division decode (what the packing commit left in
11483    // place) vs the fused sign-LUT matvec. Both single-threaded, same bytes.
11484    //   cargo test -p cortiq-engine q1t_matvec_speed -- --ignored --nocapture
11485    #[test]
11486    #[ignore]
11487    fn q1t_matvec_speed() {
11488        use cortiq_core::quant::{Q1T_TILE, f32_to_f16, q1t_code, q1t_pack};
11489        use std::time::Instant;
11490        let (rows, cols) = (8192usize, 4096usize); // FFN-sized
11491        let gpr = cols / GROUP_SIZE;
11492        let mut bytes = Vec::with_capacity(rows * gpr * Q1T_TILE + 16);
11493        for r in 0..rows {
11494            for g in 0..gpr {
11495                let s = 0.1 + ((r + g) % 7) as f32 * 0.01;
11496                bytes.extend_from_slice(&f32_to_f16(s).to_le_bytes());
11497                let mut c = [0u8; 7];
11498                for k in 0..GROUP_SIZE {
11499                    q1t_pack(&mut c, k, ((k * 7 + r + g) % 3) as u8);
11500                }
11501                bytes.extend_from_slice(&c);
11502            }
11503        }
11504        let (n, stride) = (rows * cols, 40usize); // ~2.5% outliers, per-row overlay
11505        let mut row_ptr = vec![0u32; rows + 1];
11506        let mut idx = 0usize;
11507        while idx < n {
11508            row_ptr[idx / cols + 1] += 1;
11509            idx += stride;
11510        }
11511        for r in 0..rows {
11512            row_ptr[r + 1] += row_ptr[r];
11513        }
11514        for &p in &row_ptr {
11515            bytes.extend_from_slice(&p.to_le_bytes());
11516        }
11517        let mut idx = 0usize;
11518        while idx < n {
11519            bytes.extend_from_slice(&((idx % cols) as u16).to_le_bytes());
11520            bytes.extend_from_slice(&f32_to_f16((idx % 13) as f32 * 0.1 - 0.6).to_le_bytes());
11521            idx += stride;
11522        }
11523        // On-grid ±1 so the fast path's int8 SDOT is exact vs the f32 "slow"
11524        // reference (the A/B is a timing check; values must still agree).
11525        let x: Vec<f32> = (0..cols)
11526            .map(|j| if j % 3 == 0 { 1.0 } else { -1.0 })
11527            .collect();
11528        let (rp_off, ent_off, has_ov) = q1t_overlay(&bytes, rows * gpr * Q1T_TILE, rows);
11529
11530        // "before": base-3 division decode into a buffer, then dot.
11531        let slow = |out: &mut [f32]| {
11532            let mut buf = vec![0f32; cols];
11533            for r in 0..rows {
11534                for g in 0..gpr {
11535                    let off = (r * gpr + g) * Q1T_TILE;
11536                    let s = f16_to_f32(u16::from_le_bytes([bytes[off], bytes[off + 1]]));
11537                    let codes = &bytes[off + 2..off + Q1T_TILE];
11538                    for k in 0..GROUP_SIZE {
11539                        buf[g * GROUP_SIZE + k] = match q1t_code(codes, k) {
11540                            1 => s,
11541                            2 => -s,
11542                            _ => 0.0,
11543                        };
11544                    }
11545                }
11546                out[r] = q1t_row_outlier_correction(&bytes, r, rp_off, ent_off, has_ov, &x)
11547                    + (0..cols).map(|j| buf[j] * x[j]).sum::<f32>();
11548            }
11549        };
11550        let iters = 5;
11551        let mut a = vec![0f32; rows];
11552        slow(&mut a); // warm
11553        let t = Instant::now();
11554        for _ in 0..iters {
11555            slow(&mut a);
11556        }
11557        let slow_ms = t.elapsed().as_secs_f64() * 1e3 / iters as f64;
11558
11559        let mut b = vec![0f32; rows];
11560        q1t_matvec(&bytes, &x, rows, cols, &mut b, None); // warm
11561        let t = Instant::now();
11562        for _ in 0..iters {
11563            q1t_matvec(&bytes, &x, rows, cols, &mut b, None);
11564        }
11565        let fast_ms = t.elapsed().as_secs_f64() * 1e3 / iters as f64;
11566
11567        for r in 0..rows {
11568            assert!((a[r] - b[r]).abs() < 1e-2, "mismatch row {r}");
11569        }
11570        println!(
11571            "q1t matvec {rows}x{cols} (1 thread): div-decode {slow_ms:.2} ms  fused-LUT {fast_ms:.2} ms  => {:.2}x",
11572            slow_ms / fast_ms
11573        );
11574    }
11575}
11576
11577#[cfg(test)]
11578mod gemm_bench {
11579    /// `cargo test -p cortiq-engine --release q4tp_matmat_throughput -- --ignored --nocapture`
11580    /// Times the batched q4tp GEMM at the shapes the image DiT runs
11581    /// (b=296 tokens, 2304 -> 9216), on synthetic bytes: no model, no
11582    /// mmap, no thermal drift over minutes — a kernel change shows up
11583    /// here in seconds where a full render hides it in noise.
11584    ///
11585    /// On macOS add `CMF_ACCEL=0`: this shape is over the 500k-cell mark
11586    /// where the matmat hands off to Accelerate's dequant sgemm, and
11587    /// without the opt-out both rows below measure the AMX, not the
11588    /// kernel under test.
11589    #[test]
11590    #[ignore]
11591    fn q4tp_matmat_throughput() {
11592        // 296 is a prompt-encode batch; the image DiT runs 2085 at
11593        // 512x512, where the activation panel stops fitting L2 and the
11594        // loop's shape starts to matter more than its instructions.
11595        let b: usize = std::env::var("CMF_BENCH_B")
11596            .ok()
11597            .and_then(|v| v.parse().ok())
11598            .unwrap_or(296);
11599        let (rows, cols) = (9216usize, 2304usize);
11600        let (_, _, _) = (rows, cols, b);
11601        let total =
11602            cortiq_core::quant::expected_nbytes(cortiq_core::TensorDtype::Q4TiledP, &[rows, cols])
11603                .unwrap();
11604        // Random nibbles are fine, but the row params are f16 (lo, step)
11605        // of a geometric ladder: garbage there gives exp2 of a huge
11606        // exponent, the scales come back inf, and the whole bench times
11607        // NaN arithmetic instead of the kernel.
11608        let (params_off, codes_off, _) = cortiq_core::quant::q4tp_sections(rows, cols);
11609        let mut bytes: Vec<u8> = (0..total).map(|i| (i * 37 % 251) as u8).collect();
11610        let lo = cortiq_core::quant::f32_to_f16(-4.0);
11611        let step = cortiq_core::quant::f32_to_f16(0.1);
11612        for r in 0..rows {
11613            let o = params_off + r * 4;
11614            bytes[o..o + 2].copy_from_slice(&lo.to_le_bytes());
11615            bytes[o + 2..o + 4].copy_from_slice(&step.to_le_bytes());
11616        }
11617        let _ = codes_off;
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        // A shared 48-core stand drifts ±25% run to run, which is wider
11624        // than any kernel change worth making. So: alternate the two
11625        // kernels inside one process and keep the BEST time for
11626        // each. Interleaving makes both see the same interference, and a
11627        // minimum is the one statistic another tenant cannot inflate.
11628        super::q4tp_matmat(&bytes, &xs, b, rows, cols, &mut out, pool.as_deref());
11629        let reps: usize = std::env::var("CMF_BENCH_REPS")
11630            .ok()
11631            .and_then(|v| v.parse().ok())
11632            .unwrap_or(10);
11633        let mut best = [f64::MAX; 2];
11634        let mut sums = [0f32; 2];
11635        for _ in 0..reps {
11636            for (k, w) in [(0usize, 1u8), (1usize, 2u8)] {
11637                super::Q4TP_ALT.store(w, std::sync::atomic::Ordering::Relaxed);
11638                let t = std::time::Instant::now();
11639                super::q4tp_matmat(&bytes, &xs, b, rows, cols, &mut out, pool.as_deref());
11640                best[k] = best[k].min(t.elapsed().as_secs_f64());
11641                sums[k] = out.iter().take(64).sum::<f32>();
11642            }
11643        }
11644        let flops = 2.0 * b as f64 * rows as f64 * cols as f64;
11645        for (k, name) in ["previous", "tuned   "].iter().enumerate() {
11646            println!(
11647                "q4tp matmat {rows}x{cols} b={b} {name}: {:.1} ms  {:.1} GFLOP/s  (checksum {:.3})",
11648                best[k] * 1e3,
11649                flops / best[k] / 1e9,
11650                sums[k]
11651            );
11652        }
11653        assert!(
11654            (sums[0] - sums[1]).abs() < 1e-2,
11655            "the tuned kernel changed the result: {} vs {}",
11656            sums[0],
11657            sums[1]
11658        );
11659    }
11660
11661    /// The blocked kernel must agree with the per-column path exactly —
11662    /// same weights, same activation split, only a different instruction
11663    /// mix. Shapes are chosen to hit the awkward cases: a column count
11664    /// that leaves an odd group (the 512-bit kernel does two at a time),
11665    /// and a batch that does not divide by four.
11666    #[test]
11667    fn q4tp_matmat_blocked_matches_scalar() {
11668        use std::sync::atomic::Ordering::Relaxed;
11669        // The last shape carries the image DiT's column count — 2304, so
11670        // 72 groups of accumulation, which is where a reordered sum can
11671        // actually drift — and runs through the thread pool, since the
11672        // blocked path splits rows across workers. Its row count stays
11673        // under 500k cells on purpose: above that, macOS diverts the whole
11674        // matmat to the Accelerate/AMX dequant sgemm and neither kernel
11675        // here would run.
11676        for &(rows, cols, b) in &[
11677            (64usize, 128usize, 7usize),
11678            (33, 96, 4),
11679            (16, 256, 9),
11680            (192, 2304, 37),
11681        ] {
11682            let total = cortiq_core::quant::expected_nbytes(
11683                cortiq_core::TensorDtype::Q4TiledP,
11684                &[rows, cols],
11685            )
11686            .unwrap();
11687            let (params_off, _, _) = cortiq_core::quant::q4tp_sections(rows, cols);
11688            let mut bytes: Vec<u8> = (0..total).map(|i| (i * 61 % 251) as u8).collect();
11689            let lo = cortiq_core::quant::f32_to_f16(-4.0);
11690            let step = cortiq_core::quant::f32_to_f16(0.1);
11691            for r in 0..rows {
11692                let o = params_off + r * 4;
11693                bytes[o..o + 2].copy_from_slice(&lo.to_le_bytes());
11694                bytes[o + 2..o + 4].copy_from_slice(&step.to_le_bytes());
11695            }
11696            let xs: Vec<f32> = (0..b * cols)
11697                .map(|i| ((i % 89) as f32 - 44.0) / 44.0)
11698                .collect();
11699            let mut got = vec![0f32; b * rows];
11700            let mut want = vec![0f32; b * rows];
11701            let gpr = cols / 32;
11702            let view = super::Q4tpView::new(&bytes, rows, cols);
11703            let pool = crate::pool::Pool::from_env();
11704            super::Q4TP_ALT.store(2, Relaxed);
11705            super::q4tp_matmat(&bytes, &xs, b, rows, cols, &mut got, pool.as_deref());
11706            super::Q4TP_ALT.store(1, Relaxed);
11707            super::q4tp_matmat(&bytes, &xs, b, rows, cols, &mut want, pool.as_deref());
11708            super::Q4TP_ALT.store(0, Relaxed);
11709            // Measured against the output's scale, not cell by cell: a
11710            // dot product of 2304 terms lands near zero wherever the row
11711            // and the activation nearly cancel, and there a per-cell
11712            // ratio reports 1e-3 for an absolute error of 5e-6 — f32's
11713            // own rounding, reordered. What must stay small is the error
11714            // relative to what the layer actually outputs.
11715            let scale = want.iter().fold(0f32, |m, v| m.max(v.abs())).max(1e-6);
11716            let (mut worst, mut at) = (0f32, 0usize);
11717            for (i, (g, w)) in got.iter().zip(&want).enumerate() {
11718                if (g - w).abs() > worst {
11719                    worst = (g - w).abs();
11720                    at = i;
11721                }
11722            }
11723            assert!(
11724                worst <= 1e-4 * scale,
11725                "{rows}x{cols} b={b}: blocked and scalar disagree by {worst:.3e} \
11726                 (scale {scale:.3e}) at cell {at}: {} vs {}",
11727                got[at],
11728                want[at]
11729            );
11730
11731            // "Same speed, no quality loss" is a claim about which answer
11732            // is RIGHT, not about which two agree. Both paths sum the same
11733            // 2304 products in different orders, so f64 decides: the
11734            // blocked kernel keeps sixteen partial sums and folds them at
11735            // the end, which is a shallower addition tree than the
11736            // per-column path's running scalar, and it must not be worse.
11737            let (mut e_blocked, mut e_scalar) = (0f64, 0f64);
11738            for bi in 0..b {
11739                let act = super::split_act(&xs[bi * cols..(bi + 1) * cols]);
11740                for r in 0..rows {
11741                    let mut sc = vec![0f32; gpr];
11742                    view.scales_into(r, gpr, &mut sc);
11743                    let mut exact = 0f64;
11744                    for j in 0..cols {
11745                        let (w, sq) = super::q4tp_outlier(view.nib, r, gpr, j, &sc);
11746                        exact += w as f64 * sq as f64 * act.xq[j] as f64;
11747                    }
11748                    exact *= act.sx as f64;
11749                    for &(j, xv) in &act.outliers {
11750                        let (w, sq) = super::q4tp_outlier(view.nib, r, gpr, j, &sc);
11751                        exact += w as f64 * sq as f64 * xv as f64;
11752                    }
11753                    let i = bi * rows + r;
11754                    e_blocked = e_blocked.max((got[i] as f64 - exact).abs());
11755                    e_scalar = e_scalar.max((want[i] as f64 - exact).abs());
11756                }
11757            }
11758            println!(
11759                "{rows}x{cols} b={b}: worst error vs f64 — blocked {e_blocked:.3e}, \
11760                 per-column {e_scalar:.3e}"
11761            );
11762            // An absolute bar, not a race between the two: at these
11763            // magnitudes both sit in f32's last bits, and on a small shape
11764            // whichever one happens to round the unluckiest cell "wins" by
11765            // a factor the next seed reverses.
11766            assert!(
11767                e_blocked <= 1e-5 * scale as f64 && e_scalar <= 1e-5 * scale as f64,
11768                "{rows}x{cols} b={b}: error against f64 too large — blocked \
11769                 {e_blocked:.3e}, per-column {e_scalar:.3e}, scale {scale:.3e}"
11770            );
11771        }
11772    }
11773
11774    /// The q4t twin of the throughput bench, same shape and rules, so the
11775    /// two quantisations' batch kernels can be read against each other.
11776    /// `cargo test -p cortiq-engine --release q4t_matmat_throughput -- --ignored --nocapture`
11777    #[test]
11778    #[ignore]
11779    fn q4t_matmat_throughput() {
11780        let (rows, cols, b) = (9216usize, 2304usize, 296usize);
11781        let total =
11782            cortiq_core::quant::expected_nbytes(cortiq_core::TensorDtype::Q4Tiled, &[rows, cols])
11783                .unwrap();
11784        // q4t carries a per-group f16 scale in the tile's first two bytes;
11785        // random bytes there decode to inf and the bench would time NaNs.
11786        let mut bytes: Vec<u8> = (0..total).map(|i| (i * 37 % 251) as u8).collect();
11787        let sc = cortiq_core::quant::f32_to_f16(0.02);
11788        for t in bytes.chunks_mut(super::Q4_TILE) {
11789            t[..2].copy_from_slice(&sc.to_le_bytes());
11790        }
11791        let xs: Vec<f32> = (0..b * cols)
11792            .map(|i| ((i % 97) as f32 - 48.0) / 48.0)
11793            .collect();
11794        let mut out = vec![0f32; b * rows];
11795        let pool = crate::pool::Pool::from_env();
11796        super::q4t_matmat(&bytes, &xs, b, rows, cols, &mut out, pool.as_deref());
11797        let reps: usize = std::env::var("CMF_BENCH_REPS")
11798            .ok()
11799            .and_then(|v| v.parse().ok())
11800            .unwrap_or(10);
11801        let mut best = f64::MAX;
11802        for _ in 0..reps {
11803            let t = std::time::Instant::now();
11804            super::q4t_matmat(&bytes, &xs, b, rows, cols, &mut out, pool.as_deref());
11805            best = best.min(t.elapsed().as_secs_f64());
11806        }
11807        let flops = 2.0 * b as f64 * rows as f64 * cols as f64;
11808        println!(
11809            "q4t matmat {rows}x{cols} b={b}: {:.1} ms  {:.1} GFLOP/s  (checksum {:.3})",
11810            best * 1e3,
11811            flops / best / 1e9,
11812            out.iter().take(64).sum::<f32>()
11813        );
11814    }
11815}