feat(moe): FUSED3=1 opt-in AVX2 expert matmul — 40% less matmul time, bit-identical output, off by default - #1024
feat(moe): FUSED3=1 opt-in AVX2 expert matmul — 40% less matmul time, bit-identical output, off by default#1024outtodata wants to merge 2 commits into
Conversation
release: v1.6.1 — the first day of real users, fixed the same day
… quant + gate/up pair Adds c/fused_simd.h with exact-integer IDOT kernels (quant_x_q8_avx2, matmul_q_idot_v2/v3, matmul_q_idot_pair_v3) and a FUSED3 env-gated dispatch at the MoE expert FFN call site. OFF by default; with the flag unset the code path is byte-for-byte today's main. Loud when on: startup banner prints fused3=%d. Bit-exactness verified by memcmp harness in c/tests/bench_fused3.c and by identical greedy token streams vs main (flag off and on).
Records the JustVugg/colibri survey: we are +125 commits, upstream +198 from the v1.5.0 merge-base; upstream has no Metal for deepseek_v4 (backend_metal_v4.mm is ours) though it does have c/backend_metal.mm for the older engines. Lists the three commits that cherry-pick cleanly and the ones that conflict. Marks B7 RESOLVED: the prerequisite implementation it demanded was built as COLI_V4_MOE_BATCHED, and measured against B7's own criterion (Metal off vs on, end-to-end) it gives 1.330x on p064, clearing the 1.15x bar. Corrects two premises B7 recorded that measurement has since falsified. Adds the deferred upstream PRs with reasons: JustVugg#1024 is AVX2-only with no NEON variant, JustVugg#934 is ARM NEON but fmt=6 E8/IQ3 rather than MXFP4, JustVugg#1017 helps the single-token decode path which our TTFT metric excludes.
|
Cross-linking a decision that affects how this PR should be described: #1044 established that olmoe's IDOT branch quantizes the activations, so it is not token-exact against the fp32 path — and the decision there is to return Two consequences for FUSED3:
Please retitle/reword along those lines and this is good to go — the engineering is sound, the isolation is exemplary (one new header, one engine, opt-in, |
- colibri.c: cache RoPE inv_freq keyed on (theta,qk) instead of rerunning half x powf per position; parallelize the prefill RoPE/KV-write loop on CPU-only builds (iterations write distinct KV rows, rope cache is _Thread_local; CUDA/Vulkan shadow-shrink keeps GPU builds serial); cache the compiled grammar per serve slot when the identical schema text is resent (grammar_reset == fresh setup end state) - deepseek_v4.c: stack routing buffers (<=512 experts, malloc fallback beyond); COLI_MODEL_DIRS multi-directory shard lookup via st_init_multi - resource_plan.py: sidecar cache for analyze_model (atomic tmp+replace writes, (size,mtime) signature self-invalidation) - download_fp8.py: --dest / $GLM_DEST instead of a hardcoded I:\ path - Makefile: opt-in LTO=1 (off by default: cross-TU inlining can change FP contraction and the token-exactness oracles are validated against the default build's regime, JustVugg#1044/JustVugg#1024) - tests: test_rope_invfreq (15.3M float pairs, 0 bit diffs), test_grammar_cache (ASan/UBSan ownership), test_analysis_cache
- colibri.c: cache RoPE inv_freq keyed on (theta,qk) instead of rerunning half x powf per position; parallelize the prefill RoPE/KV-write loop on CPU-only builds (iterations write distinct KV rows, rope cache is _Thread_local; CUDA/Vulkan shadow-shrink keeps GPU builds serial); cache the compiled grammar per serve slot when the identical schema text is resent (grammar_reset == fresh setup end state) - deepseek_v4.c: stack routing buffers (<=512 experts, malloc fallback beyond); COLI_MODEL_DIRS multi-directory shard lookup via st_init_multi - resource_plan.py: sidecar cache for analyze_model (atomic tmp+replace writes, (size,mtime) signature self-invalidation) - download_fp8.py: --dest / $GLM_DEST instead of a hardcoded I:\ path - Makefile: opt-in LTO=1 (off by default: cross-TU inlining can change FP contraction and the token-exactness oracles are validated against the default build's regime, JustVugg#1044/JustVugg#1024) - tests: test_rope_invfreq (15.3M float pairs, 0 bit diffs), test_grammar_cache (ASan/UBSan ownership), test_analysis_cache
|
Status ping: this is currently A rebase would put it back in the queue. I'm not asking as a formality — the change is still wanted, it just can't be reviewed or merged in this state. If you'd rather not carry it any further, say so and I'll close it with the reasoning recorded, so the work is findable rather than silently dropped. Either answer is fine; what I'd like to avoid is it sitting here indefinitely because nobody said anything. |
- colibri.c: cache RoPE inv_freq keyed on (theta,qk) instead of rerunning half x powf per position; parallelize the prefill RoPE/KV-write loop on CPU-only builds (iterations write distinct KV rows, rope cache is _Thread_local; CUDA/Vulkan shadow-shrink keeps GPU builds serial); cache the compiled grammar per serve slot when the identical schema text is resent (grammar_reset == fresh setup end state) - deepseek_v4.c: stack routing buffers (<=512 experts, malloc fallback beyond); COLI_MODEL_DIRS multi-directory shard lookup via st_init_multi - resource_plan.py: sidecar cache for analyze_model (atomic tmp+replace writes, (size,mtime) signature self-invalidation; the non-serializable resolved_family registry object is rebuilt from a cheap resolve_model() on cache hits, with only the scan-derived indexer flag persisted) - download_fp8.py: --dest / $GLM_DEST instead of a hardcoded I:\ path - Makefile: opt-in LTO=1 (off by default: cross-TU inlining can change FP contraction and the token-exactness oracles are validated against the default build's regime, JustVugg#1044/JustVugg#1024) - tests: test_rope_invfreq (15.3M float pairs, 0 bit diffs; built with -ffp-contract=off so the reference-vs-engine compare does not depend on contraction luck), test_grammar_cache (ASan/UBSan ownership), test_analysis_cache (temp dir resolved for macOS/Windows canonical paths)
- colibri.c: cache RoPE inv_freq keyed on (theta,qk) instead of rerunning half x powf per position; parallelize the prefill RoPE/KV-write loop on CPU-only builds (iterations write distinct KV rows, rope cache is _Thread_local; CUDA/Vulkan shadow-shrink keeps GPU builds serial); cache the compiled grammar per serve slot when the identical schema text is resent (grammar_reset == fresh setup end state) - deepseek_v4.c: stack routing buffers (<=512 experts, malloc fallback beyond); COLI_MODEL_DIRS multi-directory shard lookup via st_init_multi - resource_plan.py: sidecar cache for analyze_model (atomic tmp+replace writes, (size,mtime) signature self-invalidation; the non-serializable resolved_family registry object is rebuilt from a cheap resolve_model() on cache hits, with only the scan-derived indexer flag persisted) - download_fp8.py: --dest / $GLM_DEST instead of a hardcoded I:\ path - Makefile: opt-in LTO=1 (off by default: cross-TU inlining can change FP contraction and the token-exactness oracles are validated against the default build's regime, JustVugg#1044/JustVugg#1024) - tests: test_rope_invfreq (15.3M float pairs, 0 bit diffs; built with -ffp-contract=off so the reference-vs-engine compare does not depend on contraction luck), test_grammar_cache (ASan/UBSan ownership), test_analysis_cache (temp dir resolved for macOS/Windows canonical paths)
|
This PR just became more strategic than when it was opened, and we would like to get it in — could you rebase onto current Context: we have been auditing the expert hot path across engines, and the audit's conclusion matches what FUSED3 does — the win is fusing gate+up+silu into fewer passes while keeping bit-identical output. Your measured "40% less matmul time, bit-identical, off by default" on olmoe is exactly the shape we want to replicate on the GLM engine next, so having yours land first gives us a reviewed reference in-tree. Two things while you rebase, if you don't mind:
|
2 participants
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(envFUSED3=1, default off): (1) AVX2-vectorized activation quantization replacing the scalar tail ofquant_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)quant_x_q8quant_x_q8_avx2: vector abs-max + quant, exact same round-to-nearest-even semanticsmatmul_q_idot_v2matmul_q_idot_v3: 4 independent int32 accumulators +_mm_prefetchof the next weight rowmatmul_q_idot_pair_v3: single pass, both outputs, shared activation cache trafficDispatch is a single
if (g_fused3)inolmoe.cat the MoE FFN call site; nothing else is touched.Measured (bench_fused3_inproc, in-process, isolated kernels; medians of ≥10)
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
FUSED3defaults to 0; when set, startup printsMODE FUSED3 ...banner plus per-layer[PROF]counters already present in the build.What this is NOT
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
#if defined(__AVX2__)around v3, scalar/v2 fallback otherwise).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
c/tests/bench_fused3.c(v3_quant_bitexact=yes,v3_output_bitidentical_v2=yeson this branch).