Skip to content

kernels: AVX2 multi-token i-quant row kernels (an AVX2-only CPU decoded every i-quant row once per token) - #43

Merged
Niko1221 merged 3 commits into
Niko1221:mainfrom
pipeob0:gh/avx2-kernels
Sep 28, 2026
Merged

Niko1221 merged 3 commits into
Niko1221:mainfrom
pipeob0:gh/avx2-kernels

Conversation

@pipeob0

@pipeob0 pipeob0 commented Sep 28, 2026

Copy link
Copy Markdown
Contributor

Why

An AVX2-only CPU (Ryzen 5000 / any X3D, Ice Lake, Rocket Lake) runs a native pack, but every i-quant
gate/up row goes through ggml's ggml_vec_dot_iq*_q8_K, which decodes the row once per token. With
--spec 4 a verify window is 3-4 tokens over the same expert rows, so the decode ran 3-4 times for
nothing. iq_avx512.cpp covers 512-bit machines and q2_avx2.cpp covers Q2_0; the i-quants had no AVX2
path at all.

This is also the concrete answer to #34 ("among AVX2, AVX512 and AVX-VNNI, which helps?"): on a native
pack AVX2 is enough, but only with these kernels — otherwise AVX2 leaves ~15-20% of the expert pool on the
table.

What

src/kernels/cpu/iq_avx2.cpp — multi-token row kernels, one decode per row for the whole window:

format rows activations
IQ2_XXS, IQ2_XS, IQ3_XXS, IQ3_S, IQ2_S gate/up Q8_K
IQ4_NL down Q8_0

Per token in the window: load, _mm256_sign_epi8, maddubs, madd, add. That is the whole per-token cost.

For IQ2_XS the decode follows ggml's own AVX2 vec_dot rather than a table lookup per sign group: the
7-bit sign indices of all 16 u16 of a 128-value group become sign vectors with two shifts, a xor, one
pshufb against ggml's bit_helper parity table and an or; the four 32-value sign vectors come out of
shuffles of that one register; the 8 scale bytes become all 16 half-scales with ggml's unpack trick plus the
k_shuffle mask. The grid lookups stay scalar (ggml does the same). Alternating accumulators break the
add chain.

even_signs is the keven_signs_q2xs equivalent — that table is static in arch/x86/quants.c so it
cannot be reused from here; it is built once at load time from ksigns_iq2xs, so it cannot drift from
the shared table (checked: bit 7 of ksigns_iq2xs[i] is popcount(i) & 1, which is exactly what
bit_helper reconstructs, so the u64 table and the parity trick agree).

Dispatch is in native_expert.cpp next to the existing AVX-512 one, with STRATA_NO_IQ256 (gate/up) and
STRATA_NO_IQ4NL (down) as per-format fallbacks to ggml, and no AVX-512 gate — so the path is exercisable
on any AVX2 machine. generate.cpp now prints which kernel is live, because the old line claimed AVX2
kernels even when the code had fallen back to vec_dot.

Correctness

native_expert_parity against ggml's vec_dot, every layer of the pack: rel 3.1e-08 … 3.4e-08, 0
failures
. Full-expert rel vs the GPU (1.2e-2 cpu / 1.1e-2 gpu) is unchanged from before this PR, i.e. the
activation quantization still dominates and the kernels are equivalent to the reference.

Measured

Machine: Ryzen 7 5700X3D (Zen 3, AVX2 only, 8 cores), 64 GB DDR4-2666, RTX 5060 Ti 16 GB at PCIe 4.0 x8,
native swift-iq3_xxs pack, --spec 4.

Microbench, one expert's rows, single thread, µs (3 tokens = one verify window):

rows ggml vec_dot this PR, 1 token this PR, 3 tokens
IQ2_XS gate+up 3 × 259 = 777 293 416
IQ2_XXS gate+up 3 × 260 = 780 246 337
IQ4_NL down 248 — 174 (5.3 GB/s vs 3.4)

End to end, greedy fixed-prompt runs, median of 6 interleaved old/new pairs, metric ms/round (tok/s is
not usable here: draft acceptance moves the round count by ±5% between runs):

before after
gate/up phase 15.163 ms/round 13.843 (-8.7%)
rows phases throughput 22.7 GB/s 23.8 (+4.8%)
whole round 50.89 ms/round 48.52 (-4.7%)

5 of the 6 pairs favoured the new kernels; the outlier is a run that routed the most experts of any run in
the set (5.67 vs 5.10-5.53 distinct per layer). The down phase, which is byte-identical code in both
builds, also improved (-6%): a less CPU-saturated pool. On this box the whole engine went from 43-46 tok/s
as shipped to ~59 tok/s in a direct run.

Honest limits

  • The gate/up phase is now noticeably closer to bandwidth-bound (23.8 of the ~35 GB/s this memory can
    stream), so further decode work will pay less than these numbers suggest; the next lever there is
    streaming loads, not more unrolling.
  • IQ3_S is still slower than ggml at 1 token (506 vs 426 µs) and 2.1× faster at 3: it wins only inside a
    verify window. A --spec 1 / single-token decode run should not take this path, which is why the
    dispatch keeps the token-count gate.
  • Measured on Zen 3 (AVX2, no VNNI). No timing claims about Alder Lake or newer Intel.

Also in this PR: the parity tool was unusable on an AVX2-only CPU

native_expert_parity called the AVX-512 row kernels and the AVX-512 quantizer unconditionally, so on
exactly the machine whose kernels you would want to validate it died with SIGILL before printing a single
line. Both calls are now behind cpu_avx512_ok(). The tool also gained a block that runs these AVX2 kernels
against ggml's vec_dot on Q8_0 activations and times them, so a kernel change can be validated and
measured without starting an engine.

The native pack's i-quant gate/up rows fell back to ggml's single-token vec_dot on a
CPU without AVX-512 (Ryzen 5000 series, Ice Lake, ...): the weights were decoded once
per token, and the 5700X3D's pool was decode bound.

iq_avx2.cpp adds multi-token row kernels for IQ2_XXS/IQ2_XS/IQ3_XXS/IQ3_S/IQ2_S and
IQ4_NL down rows, and for IQ2_XS the vectorized sign decode of ggml's AVX2 vec_dot
(the 8th sign bit by parity through the bit_helper pshufb, the sign vectors out of
shuffles of one 32-byte load, the 8 scale bytes through the unpack trick), with an
even_signs table equivalent to keven_signs_q2xs.  A row is decoded once and every
token in the verify window only pays a load, a sign, a maddubs, a madd and an add.

Measured on a 5700X3D + RTX 5060 Ti, Swift native pack, greedy, 300 tokens, medians of
interleaved A/B pairs: gate/up 15.16 -> 13.84 ms/round, rows 22.7 -> 23.8 GB/s,
48.5 ms/round vs 50.9 (-4.7%).  Single thread, one expert's gate+up: IQ2_XS 3 tokens
536 -> 416 us, IQ2_XXS 390 -> 337 us.  IQ4_NL down rows 248 -> 174 us for 3 tokens.
native_expert_parity: rel 3.1e-08 vs ggml's vec_dot, 0 failures.

STRATA_NO_IQ256 and STRATA_NO_IQ4NL fall back to ggml per format.
The tool called the AVX-512 row kernels and the AVX-512 quantizer unconditionally, so on
an AVX2-only CPU it died with SIGILL before printing anything - the one tool you want on
that machine.  Both calls are now behind cpu_avx512_ok().

Adds a block that runs the new AVX2 kernels directly against ggml's vec_dot on Q8_0
activations (the IQ4_NL down rows' real input) and times them, so a kernel change can be
validated and measured without the engine.
The notice printed 'the expert kernels run on AVX2' whether or not the AVX2 i-quant row
kernels were used, and with STRATA_NO_IQ256 set they are not: the run is on ggml's vec_dot.
Name the kernel, and name the switch that changed it.
@Niko1221

Copy link
Copy Markdown
Owner

Merged, thank you - a careful piece of work. Checked here on a Ryzen 7600 (Zen 4), IQ3_S:

  • native_expert_parity: 0 failures; the AVX2 gate/up rows 3.0-3.4e-08 and the IQ4_NL down rows ~1e-07 from ggml's vec_dot on every layer tested.
  • End to end (8 requests per arm, same answers in every arm): AVX-512 path off, your kernels vs ggml: 44.0 vs 42.8 tok/s mean (46.0 vs 43.8 median). On AVX-512 the gate/up path is unchanged; the IQ4_NL down kernel now runs there too (128 vs 178 us per 3-token call in the parity tool), neutral end to end on this 6-core CPU. Two small follow-ups if you like: the comment above the dispatch still says STRATA_NO_IQ512 falls back to ggml (it now falls to your AVX2 kernels; ggml needs both switches), and the parity tool could count a mismatch of the new kernels as a failure instead of only printing it.

@pipeob0
pipeob0 deleted the gh/avx2-kernels branch October 1, 2026 20:43
MistVVK pushed a commit to MistVVK/XeStrata that referenced this pull request Oct 3, 2026
Hugua700 pushed a commit to Hugua700/Strata that referenced this pull request Oct 6, 2026
anon761 pushed a commit to anon761/Strata-upstream that referenced this pull request Oct 6, 2026
sskver pushed a commit to sskver/Strata that referenced this pull request Oct 7, 2026
…cache

Fix prefill cache handoff between agentic tool turns
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants