i2_s AVX2 GEMM folds the int16 accumulator 10x more often than its own bound requires (7.5% measured)
Nobody has claimed this yet.
- Dominant language
- C++
- Stars
- 40.3k
- Forks
- 3.7k
- PR merge metrics
- No merged PRs in 30d
Description
Summary
tinyBLAS_I2S_AVX::gemm folds its int16 accumulator into int32 once per
128-weight block. Its own operands license once per ten. Deferring the fold
is worth a measured 7.5 % on the dispatched path, with perplexity unchanged.
This is the path that actually runs: on a default build the per-row vec_dot is
never called (calls=0 sgemm=52068 from the kernel's own counters), so the
saving is on the hot path rather than a fallback.
The bound, from the kernel's own operands
ggml/src/ggml-cpu/llamafile/sgemm.cpp, non-VNNI branch:
__m256i d0 = _mm256_maddubs_epi16(a_vals[r][0], b0);
__m256i d1 = _mm256_maddubs_epi16(a_vals[r][1], b1);
__m256i d2 = _mm256_maddubs_epi16(a_vals[r][2], b2);
__m256i d3 = _mm256_maddubs_epi16(a_vals[r][3], b3);
__m256i sum = _mm256_add_epi16(_mm256_add_epi16(d0, d1),
_mm256_add_epi16(d2, d3));
acc[r][c] = _mm256_add_epi32(acc[r][c],
_mm256_madd_epi16(sum, one16)); // every block
Codes are masked to 0..3 (the lut comment a few lines above says so) and
activations are int8, so one vpmaddubsw lane takes two products of at most
3 * 127:
per lane, per plane 2 * 3 * 127 = 762
four planes per block 4 * 762 = 3048
safe fold interval floor(32767 / 3048) = 10 blocks
With the codes an encoder actually emits (0..2; code 3 measured absent in
2,084,044,800 weights of BitNet-b1.58-2B-4T) the bound is 16. Ten is the
conservative figure and is what I measured against.
Measurement
Ryzen 5 3600 (Zen 2, AVX2, no VNNI), same binary, same session, mean of three:
folding every block 10,254,922,296 cycles
folding every ten 9,488,166,828 cycles
difference 7.5 %
llama-perplexity over the same corpus is identical between the two builds, so
this is a cost and not a trade.
Scope
- Non-VNNI only. The
__AVXVNNI__/__AVX512VNNI__branch above uses
_mm256_dpbusd_epi32, which accumulates straight toint32; there is no fold
to amortise and nothing to gain there. - A second site has the same shape and I have not benchmarked it:
ggml/src/ggml-cpu/ggml-cpu-i2s.c, inggml_gemm_i2_i8_s, builds the same
sum16from fourmaddubsand folds it per block. A fix should probably
cover both.
Suggested shape of a fix
Hoist the madd_epi16 out of the block loop and run an int16 accumulator for
up to ten blocks, folding at the boundary and at the tail. Upstream ggml does
exactly this for its own ternary types and documents the reasoning in
ggml/src/ggml-cpu/arch/x86/quants.c — TQ2_0 carries the comment
// 16-bit sums, because 256*127 still fits. That is the idiom this kernel is
missing, and the best reference for the change.
Provenance
Found while measuring this kernel as a baseline for unrelated work. The 7.5 %
figure, the bound derivation and the calls=0 sgemm=52068 dispatch counts are
reproducible from https://github.com/purpleskulll/bitnet-t5b — results/i2s_fold_interval.txt
carries the run. I have no patch to offer that I have tested against your CI;
happy to prepare one if the direction is welcome.
I searched issues and PRs for fold, accumulator, int16, tinyBLAS_I2S,
maddubs and madd_epi16 before opening this and found nothing covering it.
#259 is the nearest neighbour and is a different kernel and a different idea.
Contributor guide
No contributing guide indexed for this repository
First steps
- Read the whole issue, then the project's contributing guide.
- Comment on the issue to say you are picking it up — it saves two people doing the same work.
- Fork the repository and make your change on a branch.
- Open a pull request that references the issue number.
Research direction
Start in ggml/src/ggml-cpu/llamafile/sgemm.cpp at the non-VNNI tinyBLAS_I2S_AVX::gemm path, then compare ggml/src/ggml-cpu/ggml-cpu-i2s.c in ggml_gemm_i2_i8_s. Read the TQ2_0 handling in ggml/src/ggml-cpu/arch/x86/quants.c and benchmark the non-VNNI path before and after. Done means both sites defer safe int16 folding without changing VNNI behavior or perplexity.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- cpp
- Domain
- performance
- Issue type
- Refactor
- Difficulty
- 4/5
- Estimated time
- 3-5 days
- Activity status
- Active
- Clarity
- Mostly clear
- Newbie friendliness
- 55/100