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
7353b06 to
f4034f4
Compare
Keep microscale accumulators in registers, specialize 64/128-value groups, and reduce AVX2 dot products without a stack transpose. Fuse Hadamard sign and normalization passes and activation prefix accumulation. Allow quantization groups independent of rotation width with per-tensor divisibility fallback. Add an optional preference for fewer K splits during fixed tuning while preserving existing quantization defaults and reference paths. Extend prefix, long-K BF16, group fallback, and fixed-output regression coverage. Update the benchmark to use valid quantization groups. Validated with the Release build, 11 kernel configurations, exact-output comparisons, multi-target x86 syntax checking, and focused AddressSanitizer with leak detection.
Add shared-weight dual-int8 activation kernels, calibrated weight interchange, packed head/body routing, focused correctness coverage, and reusable offline GPTQ calibration tools.
jan-wassenberg
left a comment
jan-wassenberg
left a comment
There was a problem hiding this comment.
Thanks for updating, and congrats on the results :D Some comments on the code:
| std::vector<float> logits; | ||
| }; | ||
|
|
||
| // Versioned binary store for root-model logits. It is uncompressed so target |
There was a problem hiding this comment.
Not sure about a binary format here, how feasible is text?
|
|
||
| } // namespace | ||
|
|
||
| uint64_t ModelComparisonFingerprint(const std::string& bytes) { |
There was a problem hiding this comment.
FNV-like is a pretty terrible hash. Consider hwy/contrib/hash/highwayhash-inl.h .
| // BF16 before the existing MatMul. This intentionally adds overhead: its | ||
| // purpose is to isolate activation-quantization error from a new kernel and | ||
| // weight quantization. | ||
| static bool ActivationI8RoundtripEnabled() { |
There was a problem hiding this comment.
I think we can remove this again.
| return !candidates_.empty(); | ||
| } | ||
| void SetCandidates(std::vector<TConfig> candidates) { | ||
| void SetCandidates(std::vector<TConfig> candidates, bool tune = true) { |
There was a problem hiding this comment.
It seems more elegant to only enable a single candidate at the point where those are produced (matmul.cc), rather than introducing a tune/fixed_ flag here.
| return; | ||
| } | ||
| #endif | ||
| MMQuantizedDot4Accumulate<GEMMA_MM_I8_BIASED_B>(di32, a, b0, b1, b2, b3, c0, |
There was a problem hiding this comment.
Let's avoid the assembly. What prevents us from just using the MMQuantizedDot4Accumulate and its SumOfMulQuadAccumulate? Note that we can make any changes to Highway required. Is the issue the lack of avxvnni?
| const auto perm = hn::IndicesFromVec(df, hn::Xor(lane, bit)); | ||
| const auto other = hn::TableLookupLanes(v, perm); | ||
| const auto upper = | ||
| hn::RebindMask(df, hn::Ne(hn::And(lane, bit), hn::Zero(du))); |
There was a problem hiding this comment.
TestBit(lane, bit) can be a bit faster.
| if (fast && (signs.size() < k || cached_hash != hash_bits)) { | ||
| signs.resize(k); | ||
| for (size_t i = 0; i < k; ++i) { | ||
| signs[i] = MMI8NegativeSign(i, hash_bits) ? 0x80000000u : 0u; |
There was a problem hiding this comment.
I'd suggest vectorizing this. Happy to advise.
i could be a u16 vector initialized to Iota(du16, 0), and incremented by Set(du16, Lanes(du16)).
To go from u16 to u32, we could perhaps just BitCast from the u16 result to u32. Or better the other way around: we have u32 indices, we temporarily drop to u16 to do the Mul, then BitCast back to u32 and shift left by 16 to put the result in the upper bit of the u32.
| return value == nullptr ? fallback : atoi(value) != 0; | ||
| } | ||
|
|
||
| static inline bool MMI8FastRotate() { |
There was a problem hiding this comment.
Any reason not to just hardcode this to true?
| // Call only after every quantized weight and calibration update is final. | ||
| static HWY_NOINLINE void MMI8PackMicroB(MatPtrT<int8_t>& data) { | ||
| const size_t k = data.Cols(); | ||
| HWY_ASSERT(data.Rows() % 8 == 0 && k % 4 == 0); |
There was a problem hiding this comment.
On the critical path, prefer HWY_DASSERT.
Also, we shouldn't allocate: we can have per-thread storage in MatMulEnv and pass the thread index indicating which one to use.
2 participants
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