Port the CPU-Max fork's EmbeddingGemma 2 CPU engine work (PR #6) without the rebrand

.github/workflows/ci.yml

@@ -67,6 +67,25 @@ jobs:
       - name: Verify AVX2 code stays in runtime-dispatched kernels
         if: matrix.platform == 'linux-x86_64'
         run: python3 scripts/check-avx2-isolation.py build/embeddinggemma
+      - name: Run the EmbeddingGemma 2 CPU kernel tests at every x86 tier
+        if: matrix.platform == 'linux-x86_64'
+        run: |
+          grep -o -w -E 'avx2|avx_vnni|avx512_vnni' /proc/cpuinfo | sort -u | tr '\n' ' '; echo
+          make build/test_gemma2_attn build/test_gemma2_vnni256
+          for tier in '' avx2 scalar; do
+            echo "== EI_GEMMA2_CPU_X86=$tier"
+            EI_GEMMA2_CPU_X86=$tier ./build/test_gemma2_attn
+            EI_GEMMA2_CPU_X86=$tier ./build/test_gemma2_vnni256 300x768x768 16x2048x512 f300x512x768 F64x768x768
+          done
+      - name: Run the EmbeddingGemma 2 CPU kernel tests (ARM)
+        if: matrix.platform == 'linux-arm64' || matrix.platform == 'darwin-arm64'
+        run: |
+          make build/test_gemma2_attn build/test_gemma2_vnni256
+          ./build/test_gemma2_attn
+          ./build/test_gemma2_vnni256 300x768x768 16x2048x512 f300x512x768 F64x768x768
+      - name: Build with clang
+        if: matrix.platform == 'linux-x86_64'
+        run: make CC=clang BUILD=build-clang build-clang/embeddinggemma build-clang/test_gemma2_attn && ./build-clang/test_gemma2_attn
       - name: Verify Windows AVX2 code stays in runtime-dispatched kernels
         if: runner.os == 'Windows'
         shell: msys2 {0}

AGENTS.md

@@ -244,6 +244,28 @@ The target is the hardware ceiling under many concurrent requests (SlimServe
   not suffix): a variant name ending in a letter dispatched half the
   simdgroups and produced NaNs before that was fixed.
 
+## EmbeddingGemma 2 CPU port lessons (2026-10-18)
+
+- **Validate a remote tree from a clean `build/`.** `rsync -a` keeps the
+  Mac's source mtimes, which were older than binaries the box had built
+  minutes earlier, so `make` reused main's baseline and a whole "validation"
+  measured the wrong binary. `rm -rf build` first; check the binary mtime.
+- **Key repack caches by the weight data pointer, not the tensor struct.**
+  Callers pass stack copies of a tensor with adjusted shapes (the vision patch
+  embedding), so a struct-address key grows one entry per HTTP thread and
+  then silently falls back to another kernel per thread.
+- **The AVX-512 boxes never run the AVX-VNNI tier.** `g2x_level()` picks
+  VNNI512 on the 9950X and the EPYC, so a 256-bit kernel for Alder/Raptor Lake
+  is dead code there unless forced (`EI_GEMMA2_CPU_X86=avxvnni`); a
+  "validated" run that does not print `vnni256=1` tested nothing.
+- **Routing media through the repacked GEMM doubles media weights in RSS**
+  unless the repack releases its source pages; measure VmRSS after one
+  request of each modality, not only after text.
+- **Model semantics live in one getter** (`ei_gemma2_official_semantics`,
+  gemma2.c) read by the CPU towers, media.c and http_media.h; the GPU vision
+  towers still read the GGUF flag, so a semantics change is CPU-complete only
+  until Metal/CUDA follow.
+
 ## EmbeddingGemma 2 scaling lessons (2026-10-10)
 
 - **One CPU wave peaks at ~64 threads and collapses past it** (256-thread pool

CONTRIBUTING.md

@@ -177,9 +177,18 @@ CPU changes must pass:
 ```sh
 make clean-cpu
 make test
+make test-gemma2-attn      # tiled attention kernel vs its scalar reference and fp64
+make test-gemma2-vnni256   # Q8_0 and float GEMM routes vs fp64 (model-free)
 make test-gemma2  # requires both new GGUF files in model/
+make test-gemma2-media     # fast CPU media towers vs the reference graphs (model files)
 ```
 
+On x86-64 run the two kernel tests at every tier (`EI_GEMMA2_CPU_X86=` unset,
+`avx2`, `scalar`, and `avxvnni` on an AVX-512 CPU, which is the only way the
+256-bit repacked GEMM runs there), and keep
+`python3 scripts/check-avx2-isolation.py build/embeddinggemma` clean: AVX code
+belongs only in the `ei_avx2_*`, `g2x_*` and `q8_tile_*` kernels.
+
 Backend changes must pass the complete suite on matching hardware:
 
 ```sh
@@ -445,8 +454,50 @@ selects the unchanged reference graph):
 - x86-64 (`src/gemma2_fast_x86.c`, runtime-dispatched, generic builds): AVX2 +
   FMA, VNNI 256-bit, and on AVX-512 VNNI CPUs a repacked 512-bit GEMM (u8
   activations with a precomputed -128*sum(w) correction).
-  `EI_GEMMA2_CPU_X86=scalar|avx2|vnni256` forces a tier.
-- `EI_GEMMA2_CPU_PROFILE=1`: per-call GEMM/row/attention wall time.
+  `EI_GEMMA2_CPU_X86=scalar|avx2|vnni256|avxvnni` forces a tier (`avxvnni`
+  runs the AVX-VNNI tier on an AVX-512 CPU, for testing the kernel below).
+- `EI_GEMMA2_CPU_VNNI256` (default on at the AVX-VNNI tier, i.e. CPUs with
+  AVX-VNNI but no AVX-512 such as Alder/Raptor Lake): `src/gemma2_vnni256.c`,
+  Q8_0 weights repacked once into 8-row groups so one vpdpbusd advances eight
+  output rows; exact integer dots, fp32 summation order only. `_SHORT=0`
+  keeps waves of at most 8 tokens on the generic tile.
+- `EI_GEMMA2_CPU_SGEMM` (default on with AVX2): `src/gemma2_sgemm_x86.c`,
+  the BF16/F16/F32 projections (per-layer input, vision patch embedding, audio
+  input projection) as fp32 panels and 6-token x 16-output register tiles
+  instead of per-call row conversion.
+- `EI_GEMMA2_CPU_FLASH` (default on with AVX2): `src/gemma2_media_x86.c`,
+  tiled online-softmax attention in the text graph over (head, query block)
+  tasks with per-query key bounds; ARM keeps the NEON row attention.
+  `EI_G2M_QB` (16..128, widths <= 128) and `-DG2M_KB=32|128` are tuning hooks.
+- `EI_GEMMA2_CPU_MEDIA_FAST` (default on): `src/gemma2_vision_fast.c` and
+  `src/gemma2_audio_fast.c`, the CPU media towers on the fast graph's GEMMs,
+  tiled attention, fused threaded row passes and persistent workspaces (the
+  vision workspace grows to the 1120-soft-token cap, about 150 MB, and is
+  kept). `0` selects the scalar reference graphs in `src/gemma2.c`.
+- `EI_GEMMA2_CPU_RELEASE_WEIGHTS` (default on, Linux and macOS): once a Q8 or
+  float weight is repacked, its GGUF pages are dropped with
+  `madvise(MADV_DONTNEED)` (a private read-only mapping, re-read from the
+  file if touched); the process keeps one copy of each repacked weight.
+- `EI_CPU_HYBRID` (default on, Linux): on CPUs that report `cpu_core` and
+  `cpu_atom` the gemma2 pool is pinned P-core primaries first, then E-cores,
+  then SMT siblings; waves under `EI_GEMMA2_CPU_HYBRID_TOKENS` (512) run on
+  the P-core threads only. No effect on other CPUs.
+- `EI_GEMMA2_CPU_YIELD` (default on, CPU backend): a media tower hands the
+  engine lock to waiting text waves between its layers, so a 1120-token image
+  no longer blocks every text request for the whole tower.
+  `EI_GEMMA2_CPU_PREPACK=1` runs one text, image and audio input before the
+  server listens so the one-time repacks are not paid by first requests.
+- `EI_GEMMA2_SEMANTICS=official`: the official processor's behaviour instead
+  of llama.cpp's (the default, which the goldens in `testdata/` use): tanh
+  GELU in the vision MLP, `log(mel + 0.001)` as the log-mel floor, no
+  `MM:SS ` text before video frames. `EI_GEMMA2_VISION_GELU=tanh|quick`,
+  `EI_GEMMA2_AUDIO_LOGMEL=hf|max` and `EI_GEMMA2_VIDEO_TIMESTAMPS=0|1`
+  override one item each. The CPU towers honour the GELU choice; Metal and
+  CUDA vision still read the GGUF's `clip.use_gelu`, so the official mode is
+  CPU-complete only.
+- `EI_GEMMA2_CPU_PROFILE=1`: per-call GEMM/row/attention wall time (text
+  graph and the fast media towers); `EI_GEMMA2_REF_PROFILE=1` the same per
+  operation class for the reference media graphs.
 - `EI_POOL_SPIN_US=N`: pool workers spin N us on a published
   generation/participant ticket before parking, so back-to-back parallel
   calls skip the condition-variable wake. The EmbeddingGemma 2 CPU pool

Makefile

@@ -108,7 +108,7 @@ SRCS_CORE := src/gguf.c
 SRCS_MODEL := $(SRCS_CORE) src/model.c
 SRCS_TOKENIZER := $(SRCS_MODEL) src/tokenizer.c
 SRCS_KERNELS := src/quants.c src/kernels.c src/kernels_avx2.c
-SRCS_ENGINE := $(SRCS_TOKENIZER) $(SRCS_KERNELS) src/parallel.c src/engine.c src/gemma2.c src/gemma2_fast.c src/gemma2_fast_x86.c src/media.c src/response_cache.c
+SRCS_ENGINE := $(SRCS_TOKENIZER) $(SRCS_KERNELS) src/parallel.c src/engine.c src/gemma2.c src/gemma2_fast.c src/gemma2_fast_x86.c src/gemma2_media_x86.c src/gemma2_vnni256.c src/gemma2_sgemm_x86.c src/gemma2_vision_fast.c src/gemma2_audio_fast.c src/media.c src/response_cache.c
 SRCS_SERVICE := src/inference_service.c
 SRCS_SERVER := src/server.c src/float_format.c
 
@@ -761,3 +761,28 @@ $(BUILD)/test_windows_paths: src/test_windows_paths.c src/windows_compat.h src/c
 test-windows-download: $(BUILD)/embeddinggemma $(BUILD)/test_downloader_fixture $(BUILD)/test_windows_paths
 	./$(BUILD)/test_windows_paths
 	python3 testdata/test_windows_download.py --binary ./$(BUILD)/embeddinggemma.exe --downloader ./$(BUILD)/test_downloader_fixture.exe
+
+# Tiled attention kernel (gemma2_media_x86.c) against its scalar reference and fp64.
+$(BUILD)/test_gemma2_attn: src/test_gemma2_attn.c $(SRCS_ENGINE) $(HDRS_NATIVE) | $(BUILD)
+	$(CC) $(CFLAGS) -o $@ src/test_gemma2_attn.c $(SRCS_ENGINE) $(LDLIBS)
+
+.PHONY: test-gemma2-attn
+test-gemma2-attn: $(BUILD)/test_gemma2_attn
+	./$(BUILD)/test_gemma2_attn
+
+# Q8_0 GEMM route exactness and throughput (EI_GEMMA2_CPU_VNNI256=0: the generic tile).
+$(BUILD)/test_gemma2_vnni256: src/test_gemma2_vnni256.c $(SRCS_ENGINE) $(HDRS_NATIVE) | $(BUILD)
+	$(CC) $(CFLAGS) -o $@ src/test_gemma2_vnni256.c $(SRCS_ENGINE) $(LDLIBS)
+
+.PHONY: test-gemma2-vnni256
+test-gemma2-vnni256: $(BUILD)/test_gemma2_vnni256
+	./$(BUILD)/test_gemma2_vnni256
+
+# Fast-vs-reference CPU media tower parity and timing.
+$(BUILD)/test_gemma2_media: src/test_gemma2_media.c $(SRCS_ENGINE) $(HDRS_NATIVE) | $(BUILD)
+	$(CC) $(CFLAGS) -o $@ src/test_gemma2_media.c $(SRCS_ENGINE) $(LDLIBS)
+
+.PHONY: test-gemma2-media
+test-gemma2-media: $(BUILD)/test_gemma2_media
+	./$(BUILD)/test_gemma2_media $(GEMMA2_MODEL) $(GEMMA2_MMPROJ) vision-synthetic
+	./$(BUILD)/test_gemma2_media $(GEMMA2_MODEL) $(GEMMA2_MMPROJ) audio-synthetic

perf/optimization_status.md

@@ -4066,3 +4066,68 @@ instead of carrying their own `G2P`, and `test_gemma2_metal_routes` picks
 this GPU's shipped default (`FUSED,SG,AFUSED` without Metal 4). Interleaved
 `bench_gemma2_metal` runs of the two binaries (300x16, 64x4, 1000x4) were
 within the host's run-to-run noise with identical minimum cosines.
+
+## 2026-10-18: CPU engine work ported from the CPU-Max fork (PR #6)
+
+Yazeedx0's fork (`embeddinggemma-cpu-max`, 32 commits on d8f9e33) carried the
+EmbeddingGemma 2 CPU work below plus a rebrand and 127k lines of reports;
+the engine was re-applied on main as a series of commits (branch
+`port/pr6-cpu-engine`), each validated on an M5 Max (clang) and a Ryzen 9950X
+(GCC 15; AVX-512 tier, so the AVX-VNNI kernel runs there only under
+`EI_GEMMA2_CPU_X86=avxvnni`, a force value added for that). Defaults are
+unchanged where the fork changed them: model semantics stay llama.cpp's
+(`EI_GEMMA2_SEMANTICS=official` is opt-in), nothing from `cpu-max/` came
+across; the fork's benchmark write-up (2.46x geomean over PyTorch FP32 /
+ONNX INT8 on an i7-14700HX) is at
+https://github.com/Yazeedx0/embeddinggemma-cpu-max.
+
+What landed, with the 9950X numbers (32 threads, `test_gemma2_vnni256` /
+`test_gemma2_media`, medians):
+
+| path | flag | result |
+|---|---|---|
+| repacked 256-bit AVX-VNNI Q8_0 GEMM | `EI_GEMMA2_CPU_VNNI256` | 2.0-2.7 TMAC/s against 0.52 for the generic AVX-VNNI tile (4.7x), exact to 5e-8 |
+| packed fp32 GEMM for BF16/F32 projections | `EI_GEMMA2_CPU_SGEMM` | 0.66-1.2 TMAC/s against 0.23 for the row kernel, 1.4e-7 |
+| tiled online-softmax attention (text) | `EI_GEMMA2_CPU_FLASH` | error vs fp64 equal to the scalar kernel's (1e-5..1e-4 err/rms) |
+| fast vision tower | `EI_GEMMA2_CPU_MEDIA_FAST` | 3.8-8.7x the reference at the AVX-512 tier, 6-17x at the AVX-VNNI tier, flattened cosine 0.9991-0.9995 |
+| fast audio tower | same | 5-11x (7.8-21x AVX-VNNI tier), 0.99988-0.99995 |
+| source-page release after repack | `EI_GEMMA2_CPU_RELEASE_WEIGHTS` | server RSS after text+image+audio 925 MB against 1536 without (main 931) |
+| hybrid P/E pool, caller pinning | `EI_CPU_HYBRID` | Linux sysfs `cpu_core`/`cpu_atom` only; not measurable on these boxes |
+| CPU media towers yield the engine lock | `EI_GEMMA2_CPU_YIELD` | short text requests during a 1,100-token image: 10-14 ms with the yield, 649 ms without, 5.2 s on main (reference towers); the image itself 889 ms against 788 |
+
+Head-of-line probe (9950X, 8 HTTP workers, unique 1536x1632 images at
+`vision_tokens` 1120 so the media cache cannot answer): a short text request
+alone takes 9.4 ms; while the image runs it takes 10-14 ms with the yield,
+649 ms without it, and 5.2 s on main, where the scalar reference tower holds
+the engine for 5.3 s. Server RSS after one text, image and audio request:
+main 931 MB (299 anonymous, 632 file); the port 925 MB at the AVX-512 tier
+(904 anonymous, 21 file) and 856 MB at the AVX-VNNI tier. The fast media
+towers route their Q8 GEMMs through the repacked kernels, which doubles the
+media weights in anonymous memory (+600 MB); releasing the GGUF source pages
+after every repack, the pre-existing 512-bit one included, gives it back
+(1536 MB with `EI_GEMMA2_CPU_RELEASE_WEIGHTS=0`).
+
+Whole-forward effect against main's binary (14 HTTP requests, text 1-1500
+words, images 48-480 px, 2-12 s audio, video, mixed): byte-identical with
+every new path off; with defaults the worst cosine is 0.99958 (12 s audio)
+and 0.99982 on text, the Q8 activation requantization amplifying
+summation-order changes. `make test-gemma2` goldens (llama.cpp outputs):
+text minimum 0.999474 (0.999641 before) at the AVX-512 tier, 0.999672 at
+the AVX-VNNI tier; vision tower 0.999106 / 0.999257 (0.999159 before);
+audio tower 0.99992. On the M5 Max the fast vision tower sits at 0.999007
+against the 0.999 gate (0.999154 with the scalar graph): the NEON route of
+the same math, kept on by default but the thinnest margin in the suite.
+
+Fixed while porting: the fork's vision patch GEMM keyed its weight cache by
+the address of a stack copy of the tensor (one 2.4 MB entry per HTTP thread,
+then a per-thread fallback kernel); the repack caches now key by the weight
+data pointer. `weight_cache_lock` stays under its architecture guard (the
+fork's sgemm hook broke Linux aarch64). `aligned_alloc` is wrapped
+(`ei_aligned_alloc`) for the UCRT. Every new AVX function carries the `g2x_`
+prefix so `scripts/check-avx2-isolation.py` holds. The audio tower clamps
+q/k/v with their own bounds instead of exiting when they differ. The
+reference-graph profiler accumulators are thread-local.
+
+Not ported: `cpu-max/` (benchmark harness, reports, uv environment), README
+rewrite, installer, `--version` / startup identity, `CPU_MAX.md`.
+

src/common.h

@@ -75,4 +75,26 @@ static inline void *ei_xrealloc(void *p, size_t n) {
     return r;
 }
 
+/* Aligned allocation for SIMD weight and activation panels (alignment a power
+ * of two, at least sizeof(void *)); free with ei_aligned_free. C11
+ * aligned_alloc wants a size that is a multiple of the alignment, so the size
+ * is rounded up; the UCRT has no aligned_alloc, so Windows uses _aligned_malloc. */
+#ifdef _WIN32
+#include <malloc.h>
+static inline void *ei_aligned_alloc(size_t align, size_t n) {
+    void *p = _aligned_malloc(n ? n : 1, align);
+    if (!p) ei_die("out of memory");
+    return p;
+}
+static inline void ei_aligned_free(void *p) { _aligned_free(p); }
+#else
+static inline void *ei_aligned_alloc(size_t align, size_t n) {
+    size_t rounded = (n + align - 1) / align * align;
+    void *p = aligned_alloc(align, rounded ? rounded : align);
+    if (!p) ei_die("out of memory");
+    return p;
+}
+static inline void ei_aligned_free(void *p) { free(p); }
+#endif
+
 #endif /* EI_COMMON_H */

src/engine.c

@@ -24,6 +24,32 @@ static void set_err(char *err, size_t err_len, const char *msg) {
     }
 }
 
+/* On hybrid CPUs (P-cores + E-cores) the gemma2 CPU pool is
+ * pinned in the order P-core primary threads, E-cores, P-core SMT siblings,
+ * and records how many leading threads sit on P-cores, so the graphs can keep
+ * small waves on the fast cores (ei_thread_pool_set_cap) and spread large ones
+ * over every core. EI_CPU_HYBRID=0, a non-hybrid CPU, or more threads than
+ * allowed CPUs keep the unpinned pool. */
+static ei_thread_pool *gemma2_cpu_pool(void) {
+    const char *flag = getenv("EI_CPU_HYBRID");
+    if (!(flag && !strcmp(flag, "0"))) {
+        int32_t n = ei_cpu_count(), performance = 0;
+        int32_t *order = ei_xmalloc((size_t)n * sizeof *order);
+        int32_t ordered = ei_cpu_hybrid_order(order, n, &performance);
+        /* The pool's own count: EI_THREADS (set by --threads), else every online CPU up to 32. */
+        const char *env = getenv("EI_THREADS");
+        int32_t threads = env && *env ? atoi(env) : (n < 32 ? n : 32);
+        if (ordered > 0 && performance > 0 && threads > 0 && threads <= ordered) {
+            ei_thread_pool *pool = ei_thread_pool_create_pinned(threads, 50, order);
+            ei_thread_pool_set_performance_threads(pool, performance < threads ? performance : threads);
+            free(order);
+            return pool;
+        }
+        free(order);
+    }
+    return ei_thread_pool_create_spin(0, 50);
+}
+
 void ei_engine_load_backend(ei_engine *e, const char *model_path, const char *backend) {
     memset(e, 0, sizeof *e);
     ei_model_load(&e->model, model_path);
@@ -58,7 +84,7 @@ void ei_engine_load_backend(ei_engine *e, const char *model_path, const char *ba
         if (!e->gemma2_metal && strcmp(requested, "auto") && strcmp(requested, "cpu"))
             ei_die("EmbeddingGemma 2 backend '%s' is not available in this build", requested);
         // Short CPU waves are wake-up bound; the gemma2 pool spins 50 us.
-        if (!e->gemma2_metal) e->thread_pool = ei_thread_pool_create_spin(0, 50);
+        if (!e->gemma2_metal) e->thread_pool = gemma2_cpu_pool();
         return;
     }
     e->short_projection_threads = 6;