feat(moe): FUSED3=1 opt-in AVX2 expert matmul — 40% less matmul time, bit-identical output, off by default

[outtodata]

FUSED3: vectorized activation quantizer + 4-way IDOT + gate/up pair matmul — 40% less matmul time, bit-identical output, off by default

TL;DR. Three opt-in kernel upgrades in c/fused_simd.h (env FUSED3=1, default off): (1) AVX2-vectorized activation quantization replacing the scalar tail of quant_x_q8, (2) a 4-accumulator IDOT matmul variant with prefetch, (3) a paired gate/up matmul that shares the activation stream. On a Core i5-7300U (AVX2, 2C/4T), time inside expert matmul drops ~40%, greedy token streams are bit-identical to master, and end-to-end throughput on our production config goes 2.08 → 2.38 tok/s. No change to file formats, no approximation anywhere: all arithmetic is still exact integer IDOT.

Setting. OLMoE-1B-7B Q8_0 snapshot, our production flags, Windows + WSL pipes (i.e. the noisy setting — numbers below are conservative; on a quiet machine the deltas are larger). MoE weights streamed from disk via the async expert pipe.

What changes (all under FUSED3, default OFF)

Kernel	Before	After
quant_x_q8	scalar loop for tail + row quant	quant_x_q8_avx2: vector abs-max + quant, exact same round-to-nearest-even semantics
matmul_q_idot_v2	2-way unroll	matmul_q_idot_v3: 4 independent int32 accumulators + _mm_prefetch of the next weight row
gate/up (proj0/proj2)	two separate matmuls, activation re-read	matmul_q_idot_pair_v3: single pass, both outputs, shared activation cache traffic

Dispatch is a single if (g_fused3) in olmoe.c at the MoE FFN call site; nothing else is touched.

Measured (bench_fused3_inproc, in-process, isolated kernels; medians of ≥10)

Metric	master	FUSED3	Δ
time inside expert matmul	baseline	−40%	matmul_ms 7664→4172 per 200-token decode (bench_fused3_tokens, 8 experts, dim 2048)
E2E (quiet machine, all our prod flags)	2.08 tok/s	2.38 tok/s	+14%

Output equivalence: md5 of the 200-token greedy stream identical with and without the flag across three independent runs (d6cb4e151fa83994d62982f3cbba4190), and identical to a separately-built reference binary. This is expected — every path is exact integer arithmetic; the flag only changes instruction scheduling, never values.

How this answers the review gates from #906

• (a) off by default, loud when on: FUSED3 defaults to 0; when set, startup prints MODE FUSED3 ... banner plus per-layer [PROF] counters already present in the build.
• (b) gate semantics exact: the flag selects between two exact implementations of the same integer kernel. No approximation, no numeric mode change. Documented as such at the dispatch site.
• (c) token-exact with flag off: when off, the code path is byte-for-byte today's master (the new kernels are not even linked into the call path). Verified by identical md5 on flag-off vs master build.

What this is NOT

• Not a quantization change: weights stay Q8_0, activations stay Q8_0. Pairs naturally with the int8 path; orthogonal to any 4/5-bit weight work.
• We also tested a 2-worker parallel expert matmul (EXPERT_WORKERS=2) in the same series: it is a NO-GO from us (pipe interference eats the gain, and it introduces thread-count-dependent scheduling). Not proposed.

Caveats

• Numbers are from one machine (i5-7300U, AVX2). No ARM/NEON variant of the new kernels — the flag should refuse or fall back on non-AVX2 builds (current code: compile-time #if defined(__AVX2__) around v3, scalar/v2 fallback otherwise).
• The +14% E2E is on a disk-bound setup; on NVMe the relative gain is smaller (matmul is a smaller slice). The −40% matmul-time figure is the portable number.

Diff is ~200 lines in c/fused_simd.h + a dispatch flag + a bench binary. Happy to open the PR if the direction looks right, or to carve it differently (e.g. three separate flags) if you prefer finer-grained gating.

Related

• Proposal: opt-in HQQ 5-bit AXIS=0 MoE weights — −25% bytes, +36% E2E on disk-bound hardware, quality within +2.4% ppl of int8 #1009 — HQQ 5-bit AXIS=0 weights RFC (orthogonal: that one is a converter-side format proposal; this PR keeps weights Q8_0).
• The bench harness for the correctness gates is included as c/tests/bench_fused3.c (v3_quant_bitexact=yes, v3_output_bitidentical_v2=yes on this branch).