cpu_neon: native W*A8 int8 GEMM for prefill, replacing dequant + BLAS sgemm (needs i8mm hardware) #500

Option 3 from the #429 investigation. Scope chosen: a shared native int8
GEMM path for all quantized formats, targeting SMMLA/i8mm hardware, with
Apple M1 explicitly out of scope.

Blocker first: no machine here can validate this

SMMLA (int8 2×2×8 matrix multiply, ARMv8.6 FEAT_I8MM) is the instruction
that makes a native int8 GEMM competitive. Neither host has it. This issue
cannot be measured on current hardware — it needs an M4/M5 Mac or an
ARMv8.6+ server. That is the first thing to resolve, before any code.

On M1/A76 the available primitive is SDOT (4 int8 MACs per lane), against
Accelerate's fp32 sgemm running on the AMX coprocessor, which is not directly
programmable and which we measure at 940–1490 GFLOP/s. A SDOT GEMM is not
expected to win there, which is why M1 Max is out of scope rather than the
target.

Why this is worth doing anyway

From the #429 work, the CPU prefill gap to the PrismML fork is ~36 % at pp256
(fork 18.75 t/s, us 13.72 best). The profile localizes it: our dequant is as
fast as theirs (1.16×), but we spend 2.69× as many samples inside BLAS for
the same matmuls, same library.

The cause is the interaction with Accelerate, not the math. cpu_neon sets
OMP_WAIT_POLICY=active on purpose (backend.c, worth 163 → 46 ms/tok at
decode), and those spinning workers deny Accelerate the cores it needs to
thread an sgemm. Measured on Bonsai-27B pp256, KMP_BLOCKTIME=0 is worth
+33 % to +60 % of prefill across every load level tested (22, 26, 31, 50, 142).

Both cheaper fixes for that failed:

• perf: match Accelerate's threading to the caller instead of oversubscribing #457 (match Accelerate's threading to the caller) — no measurable
  effect, closed as a negative result.
• Phase-scoped kmp_set_blocktime — built and measured in two variants
  (from the entering thread, and from every worker inside a parallel region).
  The hook fires, the symbol resolves, and neither changes throughput. libomp
  appears to take the blocktime at pool construction, where
  OMP_WAIT_POLICY=active fixes it, and it is not reachable at runtime.
• Making the sgemm calls bigger (the fork's own structure) is slower
  end to end — see the panel sweep in
  geistlib#429 (comment) — because the intermediate no longer fits in cache.

That leaves two routes: change the global wait policy and pay a decode price
(measurement pending), or stop handing prefill to BLAS at all. This issue
is the second. A native int8 GEMM removes the conflict structurally: one
thread pool, no fp32 intermediate, no dequant round trip through DRAM.

Scope

A shared W*A8 native GEMM path in cpu_neon for all quantized formats — not
PQ2_0 only. Today 14 kernel-table rows route to
cpu_neon_w_dequant_trampoline_mN on mac, and cpu_neon has no native
quant GEMM at all: all 7 geist_sgemm call sites go through BLAS.

cpu_x86 already established this structure and it should be mirrored rather
than reinvented:

src/backends/cpu_x86/kernel_w8a8.{h,c}            inner kernel + contract
src/backends/cpu_x86/kernel_w8a8_avx512_vnni.c    ISA specialization
src/backends/cpu_x86/kernel_w8a8_scalar.c         portable fallback
src/backends/cpu_x86/q6k_to_w8a8.{h,c}            load-time repack
                       (and the w4a8 / q4k_to_w4a8 set alongside)

w8a8_gemv and w8a8_gemm are both there, so the shape of the API, the
repack-at-resolve pattern and the scalar-fallback discipline are all settled
questions. The ARM work is the SMMLA/SDOT inner kernel plus the per-format
repackers.

Design sketch

1. kernel_w8a8_neon.c / kernel_w4a8_neon.c — SMMLA inner kernels with an
   SDOT path for hosts without i8mm (correct, not fast; the scalar fallback
   discipline cpu_x86 already follows).
2. Per-format repackers at resolve time, following q6k_to_w8a8 /
   q4k_to_w4a8. PQ2_0 already has an x8 repack (pq2_0_x8_pack) whose
   block layout may be reusable.
3. Install as linear_mN where the repack is eligible, keeping the dequant
   trampoline as the fallback — the same shape as
   install_pq2_0_x8_gemv_if_eligible.
4. Activation quantization is shared with the existing GEMV path
   (ws_quantize_act, pq2_0_prep), so the A8 side already exists.
5. The linear_pair_* slots should be filled from the start
   (geistlib#462 added them for PQ2_0) — a native GEMM has the same shared-activation
   opportunity.

Acceptance criteria

• On i8mm hardware: pp256 prefill on Bonsai-27B at or above the PrismML
  fork's, with OMP_WAIT_POLICY=active left unchanged (the point is to
  stop needing BLAS during prefill, not to trade decode away).
• Bit-exactness is not expected — this changes the arithmetic from
  dequant-to-fp32 to direct int8. Needs a tolerance gate against the existing
  W2A8/W8A8 reference models, plus test_bonsai_e2e_int greedy-token
  agreement.
• No regression on hosts without i8mm: they keep the trampoline, so the
  eligibility check has to be a hardware probe, not a build flag.
• RSS: the repack duplicates weights the way the x8 one does (8.0 → 14.4 GB
  on the 27B). A native path that reads the source layout directly would
  avoid that and is worth considering in the design.

Explicitly out of scope

• Apple M1/M2 and Cortex-A76. No i8mm; the SDOT path exists for
  correctness only. Do not expect or measure a win there.
• The decode GEMV path, which is already native NEON (SDOT) and not affected.
• The global wait-policy question, which is the other route and stands on its
  own.

Related

• geistlib#429 — the measurement chain and everything ruled out
• geistlib#457 — closed negative result (Accelerate threading)
• geistlib#462, geistlib#485 — the efficiency items, together ~1.2 %, not the gap

🤖 Generated with Claude Code