#1949 cpu(bug): AVX2 router dot keeps accumulators in registers (2.1x slower on AMD Zen 3 and Zen 5)

本文

## Title

cpu: the AVX2 router dot keeps its accumulators in registers (2.1x on AMD Zen 3 and Zen 5; no compiler does it for you)

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.

## What changed

`src/kernels/cpu/kq_avx2.cpp`, one function, `bf16_rows_dot_multi` (+22/−10), 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, the 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 expert regression:** 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.
- **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 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.