quant-matmul

W4A8, W8A8, and FP16-storage linear layers for quantized LLM inference on CPUs, aarch64 first, loadable through kernels. Int4 weights in GPTQ/AWQ-compatible group quantization run against per-token INT8 activations. The reference baseline is a dequantized fp32 product, matched within 1e-2. Companions: bitnet-cpu (ternary), cpu-attn (int8 KV attention), decode-ops (fused glue).

A 4-bit checkpoint is small enough for a single-board computer, but the Python stack has no CPU kernel that multiplies the packed nibbles directly: loaders dequantize to float first, which discards the size advantage and the speed with it. This kernel keeps the weights packed, decodes nibbles in registers against an even/odd split of the activations, and picks the widest integer dot-product instruction the CPU actually has, so a 4096-square decode layer runs in about a millisecond on a Raspberry Pi where fp32 takes forty.

Measured on the boards: 1.10 ms for a 4096-square W4A8 decode layer on a Pi 5's Cortex-A76 against 40.6 ms in fp32 (37x) with weights down from 64 MB to 8 MB, 3.44 ms for an 11008-wide FFN layer (31x), and the same kernel at 8.76 and 22.68 ms on a Pi 4's A72, within 1e-2 of a dequantized fp32 reference throughout.

Usage

import torch
from kernels import get_kernel

qm = get_kernel("phanerozoic/quant-matmul", version=1, trust_remote_code=True)

# From a GPTQ checkpoint's tensors:
packed, s, z, gs, perm = qm.from_gptq(qweight, qzeros, scales, g_idx)
y = qm.w4a8_linear(x, packed, s, z, gs, perm)            # [..., N]

# Or as a module:
layer = qm.QuantLinear.from_gptq_tensors(qweight, qzeros, scales, g_idx)
y = layer(x)

version selects the release branch; trust_remote_code is required by kernels for publishers without the trusted-publisher mark.

API

zero_offset=1 follows the AutoGPTQ storage convention; pass 0 for checkpoints without it.

Method

Per group g of group_size weights along K:

y[m, n] = sa[m] * sum_g s[n, g] * (dot_g(a_m, w_n) - z[n, g] * asum_g[m])

Group dots are exact in int32 (nibble <= 15, |a| <= 127); groups accumulate in fp32 in fixed order, so every dispatch tier produces bit-identical f32 output. Nibbles are shifted and masked in registers against a one-time even/odd split of the activations; no unpack buffer exists.

Chosen once at runtime from CPUID / HWCAP; QM_CPU_ISA demotes the selection for A/B runs and never promotes.

Measured

Raspberry Pi 5 (Cortex-A76, SDOT) and Pi 4 Model B (Cortex-A72, NEON), torch 2.13 CPU, group_size 128, against torch fp32 on the same board:

Further measurements at group_size 128 on the same boards: w4a8 1, 4096, 11008 runs 3.57 ms (Pi 5) and 10.4 ms (Pi 4); 128, 4096, 4096 runs 23.2 and 209 ms; fp16-storage 1, 4096, 4096 runs 2.80 and 13.5 ms. Decode streams packed weights at 8-12 GB/s on the Pi 5, the board's memory ceiling; M=128 runs 185 GOPS on the SDOT tier. On x86-64 (12th-gen Intel mobile, AVX-VNNI), W4A8 M=1 4096-square runs 0.16 ms (1.6x its scalar path) and the fp16-storage GEMV 0.25 ms (7x).

Correctness

Output is within 1e-2 relative of a dequantized fp32 reference across the measured shapes on both boards, which is the int8 activation quantization step. Group dots are exact in int32 and groups accumulate in a fixed order, so every dispatch tier gives bit-identical f32 output; this is verified by demoting through QM_CPU_ISA.

Requirements and limits

• K divisible by group_size; group_size a multiple of 32 (64 and 128 tested).
• bf16 or f32 activations; output follows the input dtype.
• GPTQ act-order (g_idx) is supported through an activation permutation applied per forward.
• Fast paths cover aarch64 (NEON; SDOT with dotprod) and x86-64 (AVX2; AVX-VNNI where present); anything else uses the correct scalar fallback.
• llama.cpp remains faster as a standalone runtime; this is for models composed in Python.

References

Frantar et al., "GPTQ" (2023); Lin et al., "AWQ" (2023); Arm sdot; Intel AVX-VNNI vpdpbusd.

License

Apache-2.0.