Add experimental W8A8 MatMul with optional QuaRot-style rotation

[Mikyx-1]

W8A8 Implementation

Date: 2026-08-31
Tested commit: 7353b0600ab4675f1fc863aaae0f34e90157eb5b

Executive summary

This work adds an experimental W8A8 matrix-multiplication path to gemma.cpp
and evaluates two modes:

• Naive W8A8: symmetric 8-bit weights and 8-bit activations without
  outlier mitigation.
• Rotated W8A8: the same W8A8 kernel, preceded by matching block-128
  randomized Hadamard transforms on weights and activations.

Both modes accelerate the complete 83-question repository MMLU fixture on the
tested x86 machine. Naive W8A8 is fastest: 1.432x control speed on Gemma 3 270M
and 1.561x on Gemma 3 1B. Rotated W8A8 retains most of that speed and preserves
behavior much better, reducing mean KL(control || W8A8) by 91.0% on 270M and
83.2% on 1B relative to naive W8A8. On 1B, rotated W8A8 exactly preserves the
control's 25/83 correctness outcome; naive W8A8 scores 23/83.

The largest process peak was 2786.6 MiB RSS. Available system memory remained
at least approximately 10 GiB and free disk approximately 244 GiB during
monitoring. No run crashed or approached the safety thresholds.

Implementation

Direction implemented

The implementation targets the weight-side decompression cost in the existing
matrix multiplication. The standard compressed-weight path repeatedly
converts weight tiles to BF16 before multiplication. W8A8 instead retains a
packed 8-bit weight representation and sends those bytes directly to integer
dot-product instructions.

The change has four connected parts:

1. W8A8 kernel: consumes quantized activations and weights, accumulates in
   int32, applies activation and weight scales, and writes the expected model
   output type.
2. Model routing: eligible MatMuls are intercepted before the original path
   when GEMMA_MM_I8=1; ineligible calls fall back unchanged.
3. Lazy packing: each weight tensor is decompressed, quantized once on first
   use, and retained in a process-wide cache. Activations are quantized per
   call because they change per token.
4. Optional rotation: when GEMMA_MM_I8_ROTATE=1, matching orthonormal
   block transforms are applied to activation and transposed-weight rows before
   quantization.

Loaded SFP weight row
  -> float decompression
  -> optional block-128 Rademacher + Hadamard rotation
  -> per-output-channel int8 quantization
  -> lazy packed-weight cache

Current activation row
  -> optional matching block-128 rotation
  -> per-token int8 quantization
  -> int8 dot products / int32 accumulation
  -> multiply by activation_scale * weight_scale
  -> BF16/F32 model output

The kernel reuses existing MMLoops blocking, threading, and autotuning. The
fused gated-FFN pair is supported through TwoMatMulI8, allowing both
projections to share one quantized activation.

Quantization method

Both operands use symmetric round-to-nearest int8 quantization with maximum
magnitude 127. For an activation row a and transposed weight row b:

s_a = max(abs(a)) / 127       q_a = round(a / s_a)
s_b = max(abs(b)) / 127       q_b = round(b / s_b)

C[row, column] ~= s_a * s_b * sum_k(q_a[k] * q_b[k])

• Activations use one runtime scale per row/token.
• Weights use one packing-time scale per transposed-B row, equivalent to one
  per output channel.
• Dot products accumulate in signed 32-bit integers across a K-range before a
  scaling step.
• Quantization has no zero point. It uses no calibration data, learned
  clipping, group scales, or mixed precision.

Signed and biased byte encodings

Activations are always signed int8. The weight encoding depends on the target:

• Arm NEON/SVE: signed int8 x int8, using native signed dot operations.
• x86: q_biased = q_b + 128 is stored as uint8, enabling efficient
  unsigned-by-signed VNNI-style dot products. The kernel subtracts
  128 * sum(q_a) for each K-range.

The x86 +128 is an instruction encoding, not an asymmetric quantization zero
point. Both encodings represent the same symmetric values and have separate
correctness binaries.

Rotation method

Rotated mode divides K into blocks of 128. Each block receives:

1. A deterministic Rademacher diagonal of fixed +1/-1 signs.
2. A Walsh-Hadamard transform.
3. Normalization by 1 / sqrt(128).

The same orthonormal transform R is applied to activations and weights:

(Ra)^T (Rb) = a^T R^T R b = a^T b

Thus the unquantized dot product is preserved. After quantization, isolated
large values have been spread across 128 coordinates, so one outlier is less
likely to set a coarse scale for the entire row. Rotation is used only when K
is divisible by 128.

Runtime controls

Variable	Purpose
GEMMA_MM_I8=1	Enable model-level W8A8 routing.
GEMMA_MM_I8_ROTATE=1	Enable block-128 matching rotations.
GEMMA_MM_I8_MIN_K=<n>	Keep matrices with K below n on the original path.
GEMMA_MM_I8_SKIP_ROWS=<n>	Keep matrices with N at least n on the original path.
GEMMA_MM_I8_INCLUDE=<list>	Restrict W8A8 to matching tensor names.
GEMMA_MM_I8_EXCLUDE=<list>	Exclude matching tensor names.
GEMMA_MM_I8_VERBOSE=1	Log each lazily packed tensor.

The reported runs used every eligible matrix and no routing filters.

Prototype boundaries

This is an evaluation harness, not a production weight loader:

• Released SFP weights are decompressed and then quantized again to int8. A
  production path should generate W8 from the original checkpoint.
• Original SFP weights and the int8 cache coexist, inflating measured RSS.
• The process-wide cache is not freed during execution.
• W8A8 and original MatMul currently share shape-only autotune keys.
• Rotation covers MatMul A/B operands only, not residual streams or KV cache.

What method? Based on what paper?

Core kernel

The core kernel is repository-specific work motivated by the weight
decompression bottleneck discussed in
google/gemma.cpp issue #560.
That discussion showed weight-tile decompression dominating parts of runtime
while activation conversion was comparatively small. This implementation
therefore consumes packed 8-bit weights directly, invokes Highway's
four-product integer dot primitive, and reuses gemma.cpp's loop nest.

The symmetric per-token-activation/per-output-channel-weight quantizer is a
simple calibration-free post-training W8A8 method. This report does not claim
that the core kernel reproduces one particular paper.

QuaRot-inspired rotation

The optional rotation is QuaRot-style, based on
QuaRot: Outlier-Free 4-Bit Inference in Rotated LLMs,
Ashkboos et al., NeurIPS 2024. QuaRot uses orthogonal rotations to remove hidden
state outliers without changing the exact full-precision computation, making
quantization easier. The paper also reports calibration-free 8-bit LLaMA-2
results using round-to-nearest quantization.

This implementation borrows the randomized-Hadamard outlier-spreading idea but
is not full QuaRot:

QuaRot	This implementation
Targets end-to-end weights, activations, and KV cache, including 4-bit inference.	Uses W8A8 MatMul operands only.
Applies model-wide invariance across residual, FFN, attention, and KV-cache components.	Applies local matching transforms to eligible MatMul A/B vectors.
Uses its own model conversion and quantization pipeline.	Reuses gemma.cpp and lazily quantizes released SFP weights.
Evaluated on LLaMA-family models and standard benchmarks.	Evaluated here on Gemma 3 270M/1B and an 83-question fixture.

The precise description is: a custom W8A8 kernel with optional
QuaRot-inspired block randomized-Hadamard preprocessing.

Evaluation metrics: “Accuracy is Not All You Need”

The benchmark's flips metric follows the definition proposed by
Accuracy is Not All You Need, Dutta,
Krishnan, Kwatra, and Ramjee, NeurIPS 2024. The paper argues that aggregate
accuracy can remain nearly unchanged after model compression even when many
individual answers exchange correctness, and therefore recommends reporting
flips and KL divergence as baseline-relative distance metrics.

For baseline results B and quantized results Q, this implementation uses:

Flips = count(B correct and Q incorrect)
      + count(B incorrect and Q correct)

Flip rate = Flips / number of questions

This matches the paper's definition: incorrect-to-incorrect transitions are
not counted as flips. evals/compare_mmlu.py additionally reports:

• correct_to_incorrect and incorrect_to_correct separately;
• wrong_to_wrong_changes; and
• answer_changes, which includes all three categories when the selected
  answer actually changes.

That extra answer_changes field is similar in purpose to the paper's broader
AllFlips analysis, while the report's correctness flips column is the
paper-compatible primary flips metric.

There are two scope differences from the paper:

1. The paper compares a higher-precision baseline with compressed variants;
   this experiment uses the released SFP model as its control and compares the
   additional W8A8 conversion against it.
2. The paper's KL definition averages vocabulary-distribution divergence over
   tokens of all answer options. This implementation records full-vocabulary
   KL(SFP control || W8A8) at the answer token for each MMLU question. The
   motivation is paper-based, but the KL sampling protocol is a local,
   lower-cost adaptation rather than an exact reproduction.

Result - Comparison

Evaluation scope

Configuration	Environment	Meaning
SFP control	GEMMA_MM_I8=0 GEMMA_MM_I8_ROTATE=0	Existing released-weight path.
W8A8 naive	GEMMA_MM_I8=1 GEMMA_MM_I8_ROTATE=0	W8A8 without rotation.
W8A8 rotated	GEMMA_MM_I8=1 GEMMA_MM_I8_ROTATE=1	W8A8 with block-128 rotation.

All runs used --num_threads 6 --pin 1 --verbosity 0 and all 83 questions in
gemma/evals/mmlu.json. This repository fixture is not official MMLU, which
has roughly 14,000 questions. Here one question equals 1.20 percentage points,
so one- or two-question accuracy differences are not conclusive.

Hardware specifications

Component	Specification
Architecture	x86_64
CPU	12th Gen Intel Core i5-12400F
Topology	6 physical cores, 12 logical CPUs
Relevant SIMD	AVX2 and AVX-VNNI available; tests selected AVX2
Cache	L1d 288 KiB total, L2 7.5 MiB total, L3 18 MiB
RAM	15 GiB reported by free -h
Swap	4.0 GiB; about 0.94 GiB occupied before evaluation
Filesystem	449 GiB total; about 244 GiB available during evaluation

End-to-end results use the x86 biased-weight path. Signed-kernel correctness
was tested, but this machine cannot provide Arm performance data.

Software specifications

Component	Version/configuration
OS	Ubuntu 24.04.4 LTS
Kernel	Linux 6.11.0-29-generic x86_64
Compiler	GCC/G++ 13.3.0
CMake / generator	3.28.3 / Unix Makefiles
Build type	Release
Python	3.12.9
Highway revision	9d5b1261
SentencePiece revision	9045b2f
gemma.cpp commit	7353b0600ab4675f1fc863aaae0f34e90157eb5b

Combined result

Model	Configuration	Accuracy	Wall time	Speedup	Answer changes	Correctness flips	Mean KL	P95 KL	Peak RSS
270M	SFP control	20/83 (24.10%)	46.55 s	1.000x	-	-	-	-	766.3 MiB
270M	W8A8 naive	19/83 (22.89%)	32.51 s	1.432x	13/83 (15.66%)	5/83 (6.02%)	0.544739	1.464864	1037.2 MiB
270M	W8A8 rotated	18/83 (21.69%)	35.45 s	1.313x	3/83 (3.61%)	2/83 (2.41%)	0.048856	0.141595	1039.4 MiB
1B	SFP control	25/83 (30.12%)	285.79 s	1.000x	-	-	-	-	1762.1 MiB
1B	W8A8 naive	23/83 (27.71%)	183.14 s	1.561x	17/83 (20.48%)	8/83 (9.64%)	0.327469	1.060650	2728.4 MiB
1B	W8A8 rotated	25/83 (30.12%)	191.24 s	1.494x	6/83 (7.23%)	0/83 (0.00%)	0.055162	0.270375	2786.6 MiB

Definitions:

• Answer change: selected A/B/C/D differs from control, including one wrong
  option changing to another wrong option.
• Correctness flip: correct becomes incorrect or incorrect becomes correct.
• KL: full-vocabulary KL(control || target) at the answer token, in nats.
• Wall time: complete process time, including startup and lazy packing.

Speed

Model	Configuration	Throughput gain	Wall-time reduction	Cost vs naive
270M	W8A8 naive	+43.19%	30.16%	baseline W8A8
270M	W8A8 rotated	+31.30%	23.84%	9.06% slower
1B	W8A8 naive	+56.05%	35.92%	baseline W8A8
1B	W8A8 rotated	+49.44%	33.08%	4.43% slower

Both W8A8 modes materially outperform control. Naive is fastest because it
avoids activation rotation. Rotation retains most of the gain, especially on
1B, where transform overhead is smaller relative to MatMul work.

Flips and answer changes

Model	Configuration	Changes	Correct -> incorrect	Incorrect -> correct	Wrong -> wrong	Net correct
270M	W8A8 naive	13	3	2	8	-1
270M	W8A8 rotated	3	2	0	1	-2
1B	W8A8 naive	17	5	3	9	-2
1B	W8A8 rotated	6	0	0	6	0

Rotation reduces changes from 13 to 3 on 270M and 17 to 6 on 1B. The 270M
accuracy count is an exception to the strong fidelity trend: two of its three
rotated changes are correct-to-incorrect. With only 83 samples, the difference
between 18, 19, and 20 correct should be interpreted cautiously.

Directly comparing rotated against naive gives 16 changes on 270M (4
correct-to-incorrect, 3 incorrect-to-correct, 9 wrong-to-wrong) and 17 on 1B
(3 correct-to-incorrect, 5 incorrect-to-correct, 9 wrong-to-wrong).

KL divergence

Model	Configuration	Mean	Median	P95	Maximum	Mean reduction vs naive
270M	W8A8 naive	0.544739	0.382031	1.464864	2.865242	-
270M	W8A8 rotated	0.048856	0.034517	0.141595	0.219807	91.0%
1B	W8A8 naive	0.327469	0.185496	1.060650	2.593445	-
1B	W8A8 rotated	0.055162	0.011058	0.270375	0.601561	83.2%

KL is more sensitive than the final multiple-choice label. Rotation improves
every reported KL statistic for both models. This is the strongest evidence
here that outlier spreading improves per-row int8 quantization fidelity.

RAM and disk usage

Model	Configuration	Peak RSS	Increase vs control
270M	SFP control	766.3 MiB	-
270M	W8A8 naive	1037.2 MiB	+35.36%
270M	W8A8 rotated	1039.4 MiB	+35.64%
1B	SFP control	1762.1 MiB	-
1B	W8A8 naive	2728.4 MiB	+54.84%
1B	W8A8 rotated	2786.6 MiB	+58.14%

Runs were serialized and /proc/<pid>/status, free -h, df -h, and result
size were polled. Available RAM never fell below about 10 GiB; free disk stayed
near 244 GiB. Swap use did not increase, falling from about 942 MiB to 921 MiB.
The final result directory is 167 MiB. There were no crashes or OOM events.

Higher prototype RSS is expected because loaded SFP weights and the int8 cache
coexist. A W8-native loader should replace rather than duplicate weights, so
these numbers are not a production W8 memory comparison.

Kernel correctness

Encoding	Rotation	Result
Signed int8	Off	PASS, 0 failures
Signed int8	On	PASS, 0 failures
Biased uint8	Off	PASS, 0 failures
Biased uint8	On	PASS, 0 failures

The rotation-invariance test reported relative dot-product error 1.270e-07.
This proves arithmetic coverage for both encodings on this host, not Arm
end-to-end performance.

Which part is faster? Why is it faster?

Measured ordering

W8A8 naive  >  W8A8 rotated  >  SFP control
   fastest                              slowest

Naive W8A8 is 43.19% higher throughput on 270M and 56.05% on 1B. Rotated W8A8
remains 31.30% and 49.44% higher. Rotation costs 9.06% over naive on 270M and
4.43% on 1B.

Why W8A8 beats control

1. Repeated BF16 weight-tile decompression is removed. Packed W8 weights are
   consumed directly.
2. The kernel moves fewer weight bytes. Int8 values reduce cache and memory
   traffic relative to decompressed BF16.
3. Hardware integer dot instructions do more useful work per instruction.
   Highway maps the four-product helper to operations such as x86 vpdpbusd
   and Arm signed dot instructions.
4. Scaling is delayed. Symmetric quantization permits int32 accumulation
   across a K-range before applying the scale product.
5. Mature loop machinery is reused. Blocking, threading, fused MatMul, and
   autotuning remain in place.

The 1B model benefits more. A reasonable implementation-based explanation is
that larger MatMuls spend more total time in the accelerated kernel and better
amortize packing/startup. No per-layer counters were captured, so this is an
inference rather than measured cycle attribution.

Why naive beats rotated

Rotated W8A8 additionally copies/converts activation rows into temporary float
buffers, performs sign multiplication and O(K log 128) block transforms, and
rotates each weight row during first-use packing. Weight rotation is one-time;
activation rotation repeats every MatMul call. This extra fidelity has a
measurable but modest runtime cost.

Why rotation improves fidelity

One outlier can determine a row's single scale and leave too few effective int8
levels for ordinary values. The normalized randomized Hadamard transform
redistributes the outlier's energy while preserving the exact pre-quantization
dot product. The large reductions in answer changes and KL are consistent with
this mechanism.

Limitations

1. Only 83 MMLU questions were evaluated; accuracy claims are low-confidence.
2. Each timing configuration was run once, without confidence intervals.
3. Only one x86 host was measured; there is no Arm latency or power result.
4. SFP-to-int8 double quantization differs from direct-checkpoint W8.
5. RSS includes both SFP weights and an int8 cache.
6. No cross-entropy, perplexity, or generation benchmark was run.
7. No per-layer profiler or hardware counters were captured.
8. This is not full QuaRot: residual and KV-cache transformations are absent.

Reproduction

Run from the repository root.

Build

cmake --build build --target \
  gemma_mmlu matmul_i8_test matmul_i8_biased_test \
  bench_matmul_i8 bench_matmul_i8_biased -j6

Kernel tests

./build/matmul_i8_test
GEMMA_MM_I8_ROTATE=1 ./build/matmul_i8_test
./build/matmul_i8_biased_test
GEMMA_MM_I8_ROTATE=1 ./build/matmul_i8_biased_test

Each should end with PASS (0 failures).

270M comparison

python3 evals/compare_models.py \
  --config evals/w8a8_270m_comparison.json \
  --build_dir build \
  --output_dir evals/results/w8a8_comparison_20260831/270m

1B comparison

python3 evals/compare_models.py \
  --config evals/w8a8_1b_comparison.json \
  --build_dir build \
  --output_dir evals/results/w8a8_comparison_20260831/1b

The tool runs configurations serially. Control writes one full-vocabulary
reference; both W8A8 targets use it for KL.

Resource monitoring

free -h
df -h .
ps -eo pid,ppid,etime,rss,%cpu,comm,args | grep gemma_mmlu
du -sh evals/results/w8a8_comparison_20260831

On a 16 GiB-class host, conservative stops are less than 2 GiB available RAM
or less than 20 GiB disk. Do not run model configurations concurrently.

Artifacts

Reports and structured metrics

• evals/results/w8a8_comparison_20260831/REPORT.md: this report.
• 270m/comparison.md, 1b/comparison.md: generated summary tables.
• 270m/comparison.json, 1b/comparison.json: exact configuration, accuracy,
  timings, flips, KL, and peak RSS.

Raw outputs

Each model result directory contains control, naive, and rotated *.mmlu.out
and *.mmlu.err files plus *-control.root-kl.bin. The two 83 MiB reference
files dominate the 167 MiB artifact directory.

Reproduction inputs

• evals/w8a8_270m_comparison.json
• evals/w8a8_1b_comparison.json

Implementation files

• ops/matmul_i8-inl.h: quantization, rotation, kernel, storage, and entry
  points.
• ops/matmul_i8_model-inl.h: routing, lazy cache, filters, and environment.
• ops/matmul-inl.h: shared four-output integer dot helper.
• ops/ops-inl.h: hooks before the standard path.
• ops/matmul_i8_test.cc: signed/biased and rotation correctness.
• ops/bench_matmul_i8.cc: standalone kernel benchmark.
• evals/compare_models.py: end-to-end comparison runner.

[jan-wassenberg]

Excellent work, in particular I like the optimization for the x86 unsigned/signed requirement.
The results, and my understanding of activation outliers, makes clear that rotation is absolutely required.
I think we could remove the flag that enables it and just require it.

One possibly important piece we are missing is scaling. WUSH style whitening (SVD) is probably overkill here, data-free scaling seems sufficient. SmoothQuant and I think DuQuant++ use max norms. Scaling by sqrt(ratio of Wup:Wdown L2 column norms) is probably better. This just scales weights up and down to shrink the quant range (pulled into the per-block scales). Then we could also evaluate slightly smaller rotation block sizes (64), I'd be curious how that compares in terms of speed and quality.

The combination of scaling + your Hadamard corresponds to the lower bound shown by https://hal.science/hal-05711980/document - hence is highly promising. (Or at least within a log factor of the lower bound, because we do not want to mix across tokens; we are in what they call the "inner-only" setting.)

I like the form of randomization in your Hadamard. We can consider a cheaper hash function: 32-bit muls are expensive in SIMD, going down to 16-bit is faster and should be just as good.