Add W8A8 int8 MatMul with optional rotations - #1003
Open
Mikyx-1 wants to merge 4 commits into
Open
Conversation
|
Thanks for your pull request! It looks like this may be your first contribution to a Google open source project. Before we can look at your pull request, you'll need to sign a Contributor License Agreement (CLA). View this failed invocation of the CLA check for more information. For the most up to date status, view the checks section at the bottom of the pull request. |
Prototype of the int8*int8 path suggested in google#560. Unlike the existing I8Stream support, which dequantizes B to BF16 per tile via MMDecompress::DecompressB, this consumes B directly and multiplies with hn::SumOfMulQuadAccumulate (vpdpbusd on x86 VNNI, sdot/usdot on NEON, svdot on SVE), accumulating in int32. Quantization is symmetric throughout: per-token scales for A, computed on the fly in place of DecompressA, and per-output-channel scales for B baked in at pack time. No zero points, so C[r,c] = a_scale[r] * b_scale[c] * dot(qa[r], qb[c]) and int32 accumulation runs over a whole kc range before a single scaling step. MMLoops is now generic over the kernel and B type, so the int8 path reuses the existing blocking, parallelization and autotuning rather than duplicating the loop nest. The BF16 path is unchanged; matmul_test still passes on all attainable targets. On x86 the 4-way dot needs one unsigned operand. B is biased by 128 there and the 128*sum_k(qa) term is subtracted per kc range using prefix sums of the quantized A. Biasing B rather than A is what makes that per-range correction cheap, and it keeps the values written to C close to the true partial sums: correcting once over the whole K would inflate the intermediates that MMAddC accumulates through C, which is unrecoverable when C is BF16 and the weight channels are not zero-mean. matmul_i8_test is built twice, once per encoding, so the x86 path is covered on non-x86 hosts. bench_matmul_i8 reports throughput against the BF16 and SFP kernels plus accuracy against an F64 reference. matmul_i8_model-inl.h routes the model's MatMuls through the kernel behind GEMMA_MM_I8=1, for end-to-end measurement only. It quantizes lazily from whatever the weights file holds, so it stacks a second quantization on top of e.g. SFP; a production path would quantize the original checkpoint and would not share MatMulEnv's shape-only autotune keys between the two kernels. Fixes #1
Mikyx-1
force-pushed
the
w8a8-int8-matmul-dev
branch
from
August 31, 2026 15:51
7353b06 to
f4034f4
Compare
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Evaluation
The complete 83-question repository MMLU fixture was run on Gemma 3 270M IT and 1B IT. Full methodology, hardware/software details, flips, KL, RAM, and raw-artifact descriptions are in issue #1002.
Rotation reduced mean KL versus naive W8A8 by 91.0% on 270M and 83.2% on 1B while retaining most of the speedup. The fixture is small, so one- or two-question accuracy differences are not statistically conclusive.
Validation
Prototype scope
The model integration is intentionally experimental. It lazily quantizes from the loaded representation, retains a process-wide cache, and currently shares shape-only MatMul autotune keys with the original path. A production integration should load weights quantized directly from the original checkpoint and use separate tuning identities.
Closes #1002