[maxious]

Title

cpu(bug): the AVX2 dot kernels keep their accumulators in registers (2.1x router; +20-24% expert rows at 4 tokens)

Issue: none — no existing issue covers this (I searched upstream for router dot / AVX2 / spill / kq_avx2; the open CPU-pool ones, #921 #1549 #1913 #1923, are about other problems). Happy to file one and link it.

Summary

bf16_rows_dot_multi (the routing-aware prefetch's router dot, kq_avx2.cpp) held its per-token sums in __m256 acc[MAXT] indexed by a variable. With eight 256-bit accumulators live, neither GCC nor clang keeps them in registers — the array-indexed accumulator is a shape no mainstream compiler undoes by itself: GCC unrolls the column loop and leaves the whole array in memory, and clang spills it too (its object carries the same stack-slot FMA plus a per-row memset/memcpy call). Every token's FMA then read and wrote a stack slot. On a Ryzen 7 5700X3D that made the AVX2 kernel slower than the AVX1 copy written for pre-Haswell CPUs, and every AVX2 CPU calls the AVX2 one (expert_source.cpp:2271). The same effect measures on a Zen 5 Ryzen 7 9800X3D (2.06x on GCC, 3.2x on clang), so this is every AVX2 CPU on both compilers, not one GCC version's codegen on one core.

Same sums, same tokens, four named accumulators instead of the array:

Ryzen 7 5700X3D (Zen 3, AVX2, no AVX-512, no AVX-VNNI), one layer's router [512 experts x 2560 embeddings], 4 tokens, one thread pinned to core 2, router_dot_parity --cpu 2, median of 3:

	before	after
AVX2 kernel bf16_rows_dot_multi	0.400–0.401 ms	0.185–0.188 ms (2.1x faster)
AVX1 kernel bf16_rows_dot_multi_avx1 (untouched)	0.248–0.249 ms	0.248–0.252 ms
AVX2 against AVX1	1.6x slower	1.3x faster
stack slots in the inner loop	21 memory ops per iteration	0
output bits (FNV-1a over out[], nt 1/2/4/8 x cols 2560/5120)	—	identical, 8 of 8

Ryzen 7 9800X3D (Zen 5, AVX2 + AVX-512), same bench, GCC 15.2.0 (the build's default compiler on this PC), one thread pinned to core 0, median of 3:

	before	after
AVX2 kernel bf16_rows_dot_multi	0.255–0.274 ms	0.124–0.125 ms (2.06x faster)
AVX1 kernel bf16_rows_dot_multi_avx1 (untouched)	0.165–0.198 ms	0.146–0.150 ms
AVX2 against AVX1	1.5x slower	1.2x faster
router_dot_parity failures	0	0

The compiler is not the variable — the array-indexed accumulator is. Both spill before the change (9800X3D, this file's exact flags -O3 -DNDEBUG -mavx2 -mfma -mf16c, 4 tokens pinned to core 0):

compiler	before: AVX2	before: AVX1	after: AVX2
GCC 15.2.0	0.256 ms	0.171 ms	0.124 ms
clang++ 21.1.8	0.401 ms	0.372 ms	0.124 ms (3.2x faster)

clang's unpatched object shows the same shape at the instruction level — vfmadd213ps 0x60(%rsp,%rsi,1),%ymm0,%ymm1 followed by vmovaps %ymm1,0x60(%rsp,%rsi,1) — so neither compiler recovers the registers itself, and on the Zen 5 box the clang build is slower still because of the per-row memset/memcpy call. With the accumulators in named registers both compilers land on the identical 0.124 ms, which is the point: the source shape, not the compiler, was the limit. It is not a flag problem either: the build carries no -march=native (the CPU kernel files get explicit per-file ISA flags), and adding it would not change kq_avx2.cpp's codegen.

The expert-row dots: the same shape, the same fix

q5_1_dot and q8_0_dot (the Q5_1 and Q8_0 down-projection rows, kq_avx2.cpp) held their per-token sums in the same runtime-indexed __m256 acc[MAXT]. Both are now template <int NT>, dispatched 1..4 and chunked four tokens at a time past that - the same idiom the router change and kq256_gu_rows already use. Measured the same way (9800X3D, GCC 15.2.0, native_expert_parity --synthetic with ggml-quantized weights and activations, one thread on core 0), at the verify window's four tokens:

kernel	format	before	after
q5_1_dot	Q5_1 down rows [2560 x 640]	190-197 us	155-156 us (1.24x)
q8_0_dot	Q8_0 down rows [2560 x 640]	135 us	104 us (1.30x)

Sums unchanged: 0 of 36 token-sets differ in any bit against ggml's own vec_dot over 1..8 tokens, both before and after; native_expert_parity: 0 failures.

Why q4k_dot is left alone. The Q4_K gate/up dot keeps its runtime-nt array on purpose. Its token accumulators are not what spills - the unpatched function carries 8 ymm stack ops for the sums, against 24 (q5_1) and 59 (q8_0); what it spills is its weight-side unpack arrays (q4l/q4h/scale_l/scale_h), which token chunking does not touch and re-does. The same change measured +5% at four tokens and -4..-7% at six to eight, so it is a wash at the window and a small regression above it, and q4k is the gate/up format that runs most.

What changed

src/kernels/cpu/kq_avx2.cpp: bf16_rows_dot_multi (the router), and q5_1_dot + q8_0_dot (the Q5_1 and Q8_0 down-projection rows); plus one comment line in kq_avx1.cpp.

• The row loop is chunked four tokens at a time (for (int t0 = 0; t0 < nt; t0 += 4)), the same idiom kq256_gu_rows already uses in this file, and each pass keeps its sums in four named __m256 accumulators (a0..a3) instead of a variable-indexed __m256 acc[MAXT]. Four live accumulators fit in registers alongside the weight vector, so the column loop can unroll without the array going to memory — under GCC and clang alike.
• A window longer than four tokens reads a row a second time, from L1, so the extra pass costs nothing measurable at nt=8 (0.37 ms, the same as the spilled form, in a micro-benchmark of the two shapes compiled with each file's exact flags).
• #pragma GCC unroll 1 on the column loop is not part of the change: it changes nothing. GCC keeps the column unroll and the spills either way; the variable-indexed array is what forces the array into memory.

What it leaves alone. The arithmetic and the reduction. Each token's FMA chain still runs over the same columns in the same order and ends in the same hsum_float_8, so every sum is bit-identical. The AVX1 kernel's code, q4k_dot (see above), the other expert row kernels (iq_avx2.cpp, q2_avx2.cpp, kq256_gu_rows, kq256_rows) and the engine's dispatch are untouched — the single kq_avx1.cpp line is its comment only, which had said the AVX1 kernel reads its rows "exactly as the AVX2 kernel does" and is no longer true once the AVX2 kernel chunks its row loop four tokens at a time. The router stays what it was — an estimate for which experts to prefetch, not part of the model's output.

Extra Notes

How I checked:

• Both CPUs: the 5700X3D (Zen 3) table above, and the 9800X3D (Zen 5) table — 0.256 → 0.124 ms, 2.06x, from the same source build and the same commands as the reproduction below.
• Both compilers: GCC 15.2.0 and clang++ 21.1.8 spill on the unpatched source (the compiler table above, and clang's object disassembled); the change removes the array-indexed accumulator, so the codegen question disappears either way — both compilers end at the same 0.124 ms.
• Bit-exact: the whole out[] hashes the same (FNV-1a) before and after for nt 1/2/4/8 at cols 2560 and 5120, 8 of 8. router_dot_parity reports 0 failures with the AVX2 rel error against the double reference unchanged at 3.16e-07.
• Tests: ctest -R 'router_dot|iq_avx|q8k_quant|q2_bitplane' — 100% tests passed, 0 failed out of 8.
• No router regression in the expert path: the iq_avx2_parity --bench run over four format pairs x 1..8 tokens (255 MB of experts streamed from DRAM, one thread) moved within 1.6% across all 64 medians, which is noise; the expert path does not call this function.
• The expert rows: native_expert_parity --synthetic q4_K/q5_1 and q4_K/q8_0 report 0 of 36 token-sets differ in any bit and 0 failures over 1..8 tokens; iq_avx2_parity, iq_avx1_parity, q8k_quant_parity and router_dot_parity all exit 0. One thread pinned to core 0 throughout.
• Backends: CUDA and HIP do not compile this file. The SYCL port does (sycl/CMakeLists.txt:176), and strata_kernels_cpu builds there under icpx 2026.1 (exit 0, no new warnings, and its object has no spills either). The expert-row change added since has not been through that SYCL build (no icpx on the measuring machine); it is the same file and the same kind of change.

The benchmark, from a source build (CMake fetches the pinned llama.cpp commit for ggml-cpu when STRATA_GGML_DIR is not given):

cmake -S . -B build -G Ninja -DCMAKE_BUILD_TYPE=Release -DSTRATA_BUILD_TESTS=ON
cmake --build build --target router_dot_parity iq_avx2_parity

This change, one thread pinned to one core (the first table is core 2 on the 5700X3D, the second core 0 on the 9800X3D; any core works, the medians move a little):

taskset -c 2 ./build/router_dot_parity     # 5700X3D: avx2 0.19 ms, avx 0.25 ms (was 0.41 ms)
taskset -c 0 ./build/router_dot_parity     # 9800X3D: avx2 0.124 ms, avx 0.171 ms (was 0.256 ms)

Both print router_dot_parity: 0 failures. The compiler comparison, in a second build folder with the same file flags:

cmake -S . -B build-clang -G Ninja -DCMAKE_BUILD_TYPE=Release \
      -DCMAKE_C_COMPILER=clang -DCMAKE_CXX_COMPILER=clang++ -DSTRATA_BUILD_TESTS=OFF
cmake --build build-clang --target router_dot_parity
taskset -c 0 ./build-clang/router_dot_parity   # before the change: avx2 0.401 ms, avx 0.372 ms

The expert rows, for the no-regression check:

taskset -c 2 ./build/iq_avx2_parity --bench --cpu 2

--cpu N pins inside the tool as well, --reps N (default 5) sets the repetitions behind each median, --pairs iq3_s/iq4_nl,... picks the format pairs and --mb 256 the expert bytes streamed from DRAM per pair.

[ggeorgovassilis]

Tested on an RTX 4060 Laptop GPU (8 GB VRAM), Ryzen 7 7735HS, 62 GB RAM, Ubuntu 24.04 (no swap), CUDA 13.0 / driver 580.178.04.

On 8 GB VRAM the Qwen3.6 family's experts don't fit, so I exercised its q36 path (decode on the CPU pool, 0 expert slots): Qwen3.6-35B-A3B-UD-Q4_K_XL.gguf, --kv int8 --mtp 2, 32K context, through the served API.

Correctness: bit-exact — the greedy output hash on a fixed smoke prompt is unchanged before/after (sha256 f25ec8fc6412c0eb). This Q4_K pack uses q4k_dot (not the changed q5_1_dot/q8_0_dot) and a GPU-side router, so the changed kernels aren't on its decode path.

Performance: flat on this path, as expected. The 2.1x router and +20-24% expert-row gains are on the main engine's CPU kernels, which I didn't run here.

[Niko1221]

Thanks Alex. We ported both commits as they are (authorship kept). Router dot and the Q5_1/Q8_0 rows give the same bits as before in our identity gates (10/10 prompts and 4K/20K prompt dumps identical on a 3060, an A4000 and an R9700 box). On our boxes (GCC 13.3) the router dot goes 0.168 -> 0.153 ms on a 5600X, 0.159 -> 0.154 on a Zen 4, 0.216 -> 0.212 on a Zen 2 (the 2x you saw needs the spilling compiler), and Q5_1/Q8_0 rows at 4 tokens gain about 4-5%. It is default with no switch since it cannot change an answer.

[Niko1221]

Thanks for the pull request.

Your change is already in main for 0.2.0-beta, with your credit, so there is nothing left for you to do here. We're closing it because we're resetting the tracker for 0.2.0-beta: there are more than 1,300 open issues and pull requests, and we're closing everything that isn't being worked on right now.

If part of your work did not make it in and you want it considered, open a new pull request on top of main that follows CONTRIBUTING.md, or comment /reopen after fixing the title area prefix and the template.