WIP: W4A8 (Q4_0 int8-dot) path for block_quant — lane-mapped SDOT kernel + wiring

[czoli1976]

What this is

A working end-to-end W4A8 path for block_quant: imported f32 matmuls with constant weights are routed to a Q4_0 (4-bit) weight × int8 activation kernel, keeping 4-bit weight storage. The headline is a lane-mapped SDOT microkernel that reaches f32-GEMM parity on the large matmul. Opening as a draft to get your steer on productionizing it (see the question at the bottom).

Builds on #2348 (merged GEMV) and #2431 (the GEMM, still open — its commit is the base here).

Pipeline

flowchart TB
    subgraph PROTO["This PR — prototype (works end-to-end on MiniLM)"]
        direction TB
        E["f32 matmul<br/>EinSum, rank-2 const weight"]
        E -->|"block_quant_w4a8 transform"| OP
        subgraph OP["W4A8MatMul op (transformers)"]
            direction LR
            WC["Q4_0 weight<br/>4-bit, N×K"]
            AQ["activation<br/>→ int8, per-block"]
            WC --> KER
            AQ --> KER
            KER["lane-mapped SDOT kernel<br/>4 cols / tile · per-block f32 rescale<br/>no horizontal reduce"]
        end
        OP --> OUT["f32 output"]
    end

    subgraph PROD["Production — direction wanted"]
        direction TB
        E2["block-quant EinSum"]
        E2 -->|"panel_extractor: Q4_0 → lane-packed int8"| MMM
        MMM["MMM int8 path<br/>executor threading + weight pre-pack"]
        MMM --> KER2["same lane-mapped SDOT kernel"]
    end

    KER -. "reuse the kernel" .-> KER2

Loading

The kernel

Q4_0::w4a8_gemm decodes 4 weight columns per tile and packs their k-values so each by-element SDOT (sdot vD.4s, vN.16b, vM.4b[lane], via inline asm — the dotprod intrinsic is nightly-only) accumulates the four columns' dot products directly into the four accumulator lanes. That removes the per-block horizontal reduce that Q4_0's per-32-block scales would otherwise force; the per-block scale is then one vectorized vmulq over the 4 lanes. Bit-exact with the scalar/GEMV path; scalar handles the n % 4 tail and non-dotprod CPUs.

Numbers (Apple Silicon, MiniLM-intermediate 64×384×1536, prefill)

kernel	µs/iter	vs f32
scalar w4a8_gemm	1652	4.9× slower
+ explicit SDOT	1521	4.5× slower
+ lane-mapped SDOT (1 thread)	645	1.9×
+ threading	419	1.2×
f32 OptMatMul (baseline)	340	—

End-to-end MiniLM (-O): correct (cosine 0.965 vs f32), 4× smaller weights, but 1.6× slower (14.97 vs 9.21 ms) — the prototype op's per-call overhead (alloc, output copy, rayon setup, activation re-quant) is paid per matmul and dominates the many small q/k/v matmuls. The kernel is at parity; the op wrapper is the overhead.

Question for you

The kernel is fast and the per-block-scale handling works. To shed the per-op overhead and get an e2e win, the clean path looks like dropping this same kernel into the MMM framework as a panel_extractor (Q4_0 → lane-packed int8) feeding the int8 path, so the executor handles threading and weight pre-packing. Does that match how you'd want it integrated, or would you route it differently (e.g. a dedicated block-quant MMM kernel)? Happy to refactor toward whatever shape you prefer.

🍍

[czoli1976]

Follow-up measurements on the decode regime (m=1), since that's where W4A8 should win and the PR table only showed prefill.

At LLM scale (4096×4096):

	µs/iter
raw Q4_0::w4a8_gemv (decode)	801 → 1.62× faster than f32
f32 mat-vec	1295
W4A8MatMul op (decode, m=1)	1327 → parity with f32

So the kernel already beats f32 by 1.62× at decode; the prototype op's per-call overhead (output alloc, rayon setup, on-the-fly activation requant, the output reshape) erodes it back to parity. That overhead is exactly what an MMM-integrated path removes — executor threading, weight pre-pack, no per-op allocation — which is the integration question in the description.

Small models don't benefit. MiniLM (384-wide layers) loses at every sequence length end-to-end (f32/W4A8 = 0.41× at seq 1, 0.59× at seq 64): its weights fit in cache, so the matmul isn't memory-bandwidth-bound and there's no 4-bit advantage to capture. The win needs large (≳4k) layers where weight traffic dominates — i.e. real decoder LLMs, not encoders.

Net: the kernel's decode win is real and sizeable; the only thing between it and an end-to-end win is the per-op overhead the MMM path sheds.