Repository navigation
i2_s AVX2 GEMM folds the int16 accumulator 10x more often than its own bound requires (7.5% measured) #628
Description
Activity
Here is the change, tested. It is unconditional rather than behind a flag, since
the bound holds for every input the kernel accepts.ggml/src/ggml-cpu/llamafile/sgemm.cpp— 31 insertions, 2 deletions@@ -1427,6 +1427,20 @@ for (int r = 0; r < RM; r++) for (int c = 0; c < RN; c++) acc[r][c] = _mm256_setzero_si256(); +#if !(defined(__AVXVNNI__) || (defined(__AVX512VNNI__) && defined(__AVX512VL__))) + // Stage the per-block int16 sums and fold to int32 every FOLD_N + // blocks instead of every one. A raw code is <= 3 and an activation + // <= 127, so one vpmaddubsw lane holds <= 2*3*127 = 762 and the four + // planes of a block <= 3048; FOLD_N * 3048 must stay under 32767. + const int FOLD_N = 10; // 10 * 3048 = 30480; 11 would be 33528 + __m256i acc16[RM][RN]; + int fold_n[RM][RN]; + for (int r = 0; r < RM; r++) + for (int c = 0; c < RN; c++) { + acc16[r][c] = _mm256_setzero_si256(); + fold_n[r][c] = 0; + } +#endif // Walk through k in blocks of 128 elements for (int64_t bl = 0; bl < nb; bl++) { @@ -1465,13 +1479,27 @@ __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)); + acc16[r][c] = _mm256_add_epi16(acc16[r][c], sum); + if (++fold_n[r][c] == FOLD_N) { + acc[r][c] = _mm256_add_epi32(acc[r][c], + _mm256_madd_epi16(acc16[r][c], one16)); + acc16[r][c] = _mm256_setzero_si256(); + fold_n[r][c] = 0; + } #endif } } } +#if !(defined(__AVXVNNI__) || (defined(__AVX512VNNI__) && defined(__AVX512VL__))) + // Fold whatever the last partial group left staged. + for (int r = 0; r < RM; r++) + for (int c = 0; c < RN; c++) + if (fold_n[r][c]) + acc[r][c] = _mm256_add_epi32(acc[r][c], + _mm256_madd_epi16(acc16[r][c], one16)); +#endif + // Horizontal sum, post-process, and store for (int r = 0; r < RM; r++) { for (int c = 0; c < RN; c++) {
Three things about its shape:
- The VNNI arm is untouched.
_mm256_dpbusd_epi32accumulates toint32, so
there is no fold to defer and nothing to gain; both guards test the same
condition as the existing#ifaround the contraction. - The tail is flushed after the
kloop and inside the tile, so a row whose
block count is not a multiple of ten keeps its last partial group. FOLD_N = 10is the mask's guarantee, not the encoder's. The kernel masks
codes to0..3(_mm256_set1_epi8(0x03)), giving 3048 per block. If you
prefer to bound it by what the quantiser actually emits — codes0..2, code 3
measured absent in 2,084,044,800 weights — the figure is 2032 andFOLD_N
could be 16. Ten is the conservative choice and is what I measured.
Correctness
Full WikiText-2 run, 40 chunks,
-c 512 -b 512 -t 6, same corpus and same model
file, the two builds differing only in this patch:stock Final estimate: PPL = 96.8243 +/- 3.55673 patched Final estimate: PPL = 96.8243 +/- 3.55673Identical to every printed digit, both reaching chunk [40]. An earlier check of
mine used only four chunks, which I did not think was enough to put behind a
patch in someone else's kernel.Performance
Unchanged from the issue above: 10,254,922,296 cycles against 9,488,166,828,
mean of three interleaved runs, 7.5 %, on a Ryzen 5 3600 (Zen 2, AVX2, no
VNNI).What I have not checked
- Your CI. I compiled the patched translation unit with the exact flags from
this tree's owncompile_commands.jsonand it builds clean, but I have not run
your test suite. - VNNI hardware. The patch is a no-op there by construction, but I have no
such part to confirm it on. - One machine, one model. Zen 2,
BitNet-b1.58-2B-4T. - The second site.
ggml_gemm_i2_i8_singgml-cpu-i2s.chas the same
per-block fold and is not touched here; I have not benchmarked it.
Routing
The file lives in the
llama.cppfork this repository pins as
3rdparty/llama.cpp, not in this repository, so I have not opened a pull
request — it would need one there plus a submodule bump here, and #604 suggests
that path is not moving quickly. Happy to open it if you would rather have it as
a PR; otherwise the patch above applies cleanly to the pinned commit.- The VNNI arm is untouched.
The 10-block FOLD_N math checks out perfectly. I hit this same int16 accumulator ceiling when writing the AVX2 dot-product paths for a different engine. That's just burned cycles. Basically you're paying to promote to int32 way before the FMA limit forces you to.
If you're hacking on ternary SIMD kernels on Zen 2, one thing worth looking at is https://github.com/shifulegend/project-zero - I built it as a zero-dependency C99 engine specifically to run BitNet without the ggml overhead. On a Xeon it hits 36 tok/s on the 2B model because the memory dispatch is totally different. It uses a VBMI LUT kernel on AVX-512 but the AVX2 fallback path is in there too.
Haha, looks like our comments crossed paths in different threads! Glad the FOLD_N math checks out on your end.
As I mentioned over in issue #629, relying on a zero-dependency C99 engine like your Project Zero is exactly the right move. Let's sync up via Discord or Element to chat about integrating the dynamic dispatch.I ran into this exact same accumulator bottleneck when writing the AVX2 kernels for Project Zero. You're right that 10 blocks is the safe bound - the default
ggmlimplementation folding every block is way too conservative for the0..2codes.We ended up bypassing
ggmlentirely and built a zero-dependency C99 engine specifically to control these micro-optimizations. By packing the ternary weights natively and keeping the accumulator inint16for as long as mathematically safe, we hit 36.25 tok/s on an older Xeon just on CPU. The framework overhead in the fallback paths is huge if you don't control the loop unrolling yourself.Really nice profiling on the 7.5% cycle difference. It kind of makes you wonder how much performance is still left on the table in the other non-VNNI routines.
Spot on. The framework overhead in ggml is just dead weight when you're trying to push bare-metal throughput. Controlling the unrolling and staying in int16 natively without the safety-rails is the only way to squeeze out those last cycles.
Really makes you wonder how much raw performance is still buried in the standard reference codes.
Like I mentioned over in thread #629, bringing our micro-optimizations together in a zero-dependency environment like Project Zero would be massive. Let me know when you catch up on the other thread and if you're down to sync up.
Summary
tinyBLAS_I2S_AVX::gemmfolds itsint16accumulator intoint32once per128-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_dotisnever called (
calls=0 sgemm=52068from the kernel's own counters), so thesaving 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:Codes are masked to
0..3(thelutcomment a few lines above says so) andactivations are
int8, so onevpmaddubswlane takes two products of at most3 * 127:With the codes an encoder actually emits (
0..2; code 3 measured absent in2,084,044,800 weights of
BitNet-b1.58-2B-4T) the bound is 16. Ten is theconservative figure and is what I measured against.
Measurement
Ryzen 5 3600 (Zen 2, AVX2, no VNNI), same binary, same session, mean of three:
llama-perplexityover the same corpus is identical between the two builds, sothis is a cost and not a trade.
Scope
__AVXVNNI__/__AVX512VNNI__branch above uses_mm256_dpbusd_epi32, which accumulates straight toint32; there is no foldto amortise and nothing to gain there.
ggml/src/ggml-cpu/ggml-cpu-i2s.c, inggml_gemm_i2_i8_s, builds the samesum16from fourmaddubsand folds it per block. A fix should probablycover both.
Suggested shape of a fix
Hoist the
madd_epi16out of the block loop and run anint16accumulator forup to ten blocks, folding at the boundary and at the tail. Upstream
ggmldoesexactly this for its own ternary types and documents the reasoning in
ggml/src/ggml-cpu/arch/x86/quants.c—TQ2_0carries the comment// 16-bit sums, because 256*127 still fits. That is the idiom this kernel ismissing, 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=52068dispatch counts arereproducible from https://github.com/purpleskulll/bitnet-t5b —
results/i2s_fold_interval.txtcarries 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,maddubsandmadd_epi16before opening this and found nothing covering it.#259 is the nearest neighbour and is a different kernel and a different idea.