i2_s AVX2 GEMM folds the int16 accumulator 10x more often than its own bound requires (7.5% measured)

[purpleskulll]

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 to int32; 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, in ggml_gemm_i2_i8_s, builds the same
  sum16 from four maddubs and 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.

[purpleskulll]

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_epi32 accumulates to int32, so
  there is no fold to defer and nothing to gain; both guards test the same
  condition as the existing #if around the contraction.
• The tail is flushed after the k loop and inside the tile, so a row whose
  block count is not a multiple of ten keeps its last partial group.
• FOLD_N = 10 is the mask's guarantee, not the encoder's. The kernel masks
  codes to 0..3 (_mm256_set1_epi8(0x03)), giving 3048 per block. If you
  prefer to bound it by what the quantiser actually emits — codes 0..2, code 3
  measured absent in 2,084,044,800 weights — the figure is 2032 and FOLD_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.55673

Identical 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 own compile_commands.json and 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_s in ggml-cpu-i2s.c has the same
  per-block fold and is not touched here; I have not benchmarked it.

Routing

The file lives in the llama.cpp fork 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.