Repository navigation
Port the CPU-Max fork's EmbeddingGemma 2 CPU engine work (PR #6) without the rebrand - #7
Merged
Merged
Conversation
…pping registry Ported from Yazeedx0/embeddinggemma-cpu-max (PR #6) without the rebrand. No numeric change: every output is byte-identical to main. - parallel: ei_thread_pool_set_cap limits how many threads join a call (the ticket stays one 64-bit release store); ei_cpu_hybrid_order lists the allowed CPUs P-core primaries first, then E-cores, then SMT siblings (Linux sysfs cpu_core/cpu_atom; 0 elsewhere). - engine: the gemma2 CPU pool is pinned in that order on hybrid CPUs and records the P-core count (EI_CPU_HYBRID=0 keeps the unpinned pool). - gguf: a registry of open read-only file mappings (ei_gguf_in_mapping), on both the mmap and the Windows MapViewOfFile paths. - media: the Hann window as a table; log-mel frames run on the engine pool on the CPU tower path; engine sections take an optional outer lock (unused until the server installs one). - gemma2: model-semantics getter (EI_GEMMA2_SEMANTICS; the default stays llama.cpp's behaviour, "official" is opt-in), the vision GELU selector, and the CPU-yield hook the media towers will call between layers. - common: ei_aligned_alloc/ei_aligned_free (posix_memalign or _aligned_malloc) for the SIMD panels that follow. Validated: make test-gemma2 on an M5 Max and on a Ryzen 9950X (GCC 15), goldens unchanged. Co-authored-by: Eric Hartford <eric.hartford@lazarus.enterprises>
Ported from Yazeedx0/embeddinggemma-cpu-max (PR #6). gemma2_media_x86.c: AVX2+FMA tiled attention (g2m_attn) in a transposed form (Q^T packed once per task, S^T and O^T as vertical vector ops, K and V read in place), fp32 throughout with a degree-7 polynomial exp; a scalar reference for other architectures and for the test; QuickGELU, SiLU and exp vectors for the media towers that follow. Every AVX body carries the g2x_ prefix so scripts/check-avx2-isolation.py keeps passing. gemma2_fast.c: the text graph's attention runs on it over (head, query block) tasks with per-query key bounds (packed sequences and the +-512 window) on x86 with AVX2; EI_GEMMA2_CPU_FLASH=0 restores attention_rows, which ARM keeps. Waves under EI_GEMMA2_CPU_HYBRID_TOKENS (512) run on a hybrid pool's P-core threads only, and the caller is pinned to the pool's first CPU for the wave (no-ops on non-hybrid pools). test_gemma2_attn: head widths 64..512, GQA/MQA, sliding windows, packed bounds and partial blocks against the scalar kernel and an fp64 ground truth (tiled error <= 2x scalar + 1e-6 and < 1e-3); `bench N D HEADS` times the kernel. The query-bound arrays are sized by G2M_ATTN_QB_MAX so an EI_G2M_QB override cannot overflow them. Validated on a Ryzen 9950X (GCC 15, clean build): the attention test at EI_GEMMA2_CPU_X86 default/avx2/scalar, the isolation script, HTTP responses byte-identical to main with EI_GEMMA2_CPU_FLASH=0 on 14 text, image, audio, video and mixed requests, and make test-gemma2 with it on (text golden minimum 0.999485 against 0.999641 before: the Q8 activation requantization amplifies the summation-order change; worst whole-forward cosine against main 0.99965 on a 600-word text). On an M5 Max the tiled path is the scalar reference and make test-gemma2 is unchanged. Co-authored-by: Eric Hartford <eric.hartford@lazarus.enterprises>
Ported from Yazeedx0/embeddinggemma-cpu-max (PR #6). gemma2_vnni256.c: for CPUs with AVX-VNNI but no AVX-512 (Alder/Raptor Lake), each Q8_0 weight is repacked once into 8-row groups so one 32-byte load feeds eight output rows of a vpdpbusd; the +128 activation bias is cancelled by a -128 * sum(w) seed, scales apply per block with one convert, multiply and FMA, and no horizontal reduction remains. Integer dot products are exact, so results equal the reference Q8_0 dot up to fp32 summation order. The repack runs on the caller's pool; its cache is keyed by the weight data pointer (stable per mapping). EI_GEMMA2_CPU_VNNI256=0 keeps the generic tile; _SHORT=0 keeps waves of <= 8 tokens on it. gemma2_release.h: after a repack the GGUF source pages of that tensor are dropped with madvise(MADV_DONTNEED) (private read-only file mapping, so a refault re-reads the file); only pages inside a registered mapping are touched. EI_GEMMA2_CPU_RELEASE_WEIGHTS=0 keeps them. gemma2_fast_x86.c: AVX-VNNI is read from CPUID leaf 7.1 directly (clang before 19 rejects __builtin_cpu_supports("avxvnni")); EI_GEMMA2_CPU_X86=avxvnni forces the AVX-VNNI tier on an AVX-512 CPU so the kernel can be tested there. The macro-generated tiles carry the g2x_ prefix for the isolation script. gemma2_fast.c exports the Q8/float GEMM, quantize, sum-squares and GELU helpers for the test and the CPU media towers. test_gemma2_vnni256: eight vision/audio/text shapes against an fp64 scalar Q8_0 dot (relative error < 1e-5), with f/F cases for the float GEMM. Validated on a Ryzen 9950X (GCC 15, clean build, EI_GEMMA2_CPU_X86=avxvnni since its native tier is AVX-512): max relative error 5e-8 on every shape at 2.0-2.7 TMAC/s against 0.52 for the generic AVX-VNNI tile; the isolation script; HTTP responses byte-identical to main with the new paths off; make test-gemma2 under the forced tier (text golden minimum 0.999659); server RSS after text, image and audio 801 MB with page release against 944 without (main: 931). On the Mac the kernel is inert and the test exercises the NEON route. Co-authored-by: Eric Hartford <eric.hartford@lazarus.enterprises>
Ported from Yazeedx0/embeddinggemma-cpu-max (PR #6). gemma2_sgemm_x86.c: the BF16 per-layer-input projection, the F32 vision patch embedding and the audio input projection ran f_rows, which converts four weight rows per task on every call and ends each 4x4 tile in sixteen horizontal reductions. The weight is now expanded once (BF16/F16 -> fp32 is exact) into 16-column panels and a 6-token x 16-output register tile streams one panel row per k with no reductions; fp32 throughout, only the summation order changes. AVX2+FMA, runtime dispatched, g2x_ prefixed; EI_GEMMA2_CPU_SGEMM=0 restores f_rows. The hook is compiled on x86-64 only, so the weight-cache lock keeps its existing architecture guard (Linux aarch64 builds). The cache is keyed by the weight data pointer: callers pass tensor copies with adjusted shapes, so the struct address would key on a stack slot. Validated on a Ryzen 9950X (GCC 15, clean build): BF16 and F32 cases of test_gemma2_vnni256 exact to 1.4e-7 against fp64 at 0.66-1.2 TMAC/s against 0.23 for f_rows; HTTP responses byte-identical to main with the new paths off; make test-gemma2 at the native tier (text golden minimum 0.999474) and the forced AVX-VNNI tier (0.999672). Inert on the Mac. Co-authored-by: Eric Hartford <eric.hartford@lazarus.enterprises>
Ported from Yazeedx0/embeddinggemma-cpu-max (PR #6). The CPU media towers ran the scalar reference graph in gemma2.c: per-element GEMM loops, scalar attention with an allocation per job, single-threaded norms, RoPE, activations and quantization. gemma2_vision_fast.c and gemma2_audio_fast.c keep the reference math (Q8_0 activations for Q8_0 weights, fp32 attention and norms, the fp16-rounded patch input, the clipped-linear clamps, the reference tensor mapping) and change the execution: Q8 GEMMs on the fast text graph's routes, tiled attention for the 64-wide vision heads with head-major K/V planes, channel-vectorized subsampling convolutions with fused LayerNorm+ReLU streamed per task, per-token fused row passes threaded over tokens, the 3072-wide FFN in token chunks, relative-position keys computed once per projector, and persistent workspaces (the vision one grows to the 1120-soft-token cap, about 150 MB, and is kept). ei_gemma2_vision/_audio dispatch to them; EI_GEMMA2_CPU_MEDIA_FAST=0 selects the reference graphs, kept as _ref. The audio tower clamps q, k and v with their own output bounds instead of exiting when they differ, and keys its per-projector cache by a weight's data pointer rather than the projector struct address. gemma2.c: EI_GEMMA2_REF_PROFILE=1 per-operation timers for the reference graphs (thread-local accumulators); the vision MLP activation comes from ei_gemma2_vision_tanh_gelu (GGUF flag or the opt-in official semantics). http_media.h: the "MM:SS " text before video frames stays by default; EI_GEMMA2_SEMANTICS=official or EI_GEMMA2_VIDEO_TIMESTAMPS=0 drops it. test_gemma2_media: fast against reference towers on synthetic images and mel clips (flattened cosine >= 0.999, every token >= 0.95); `make test-gemma2-media` runs both. Validated: Ryzen 9950X (GCC 15, clean build) at the native, forced AVX-VNNI and scalar tiers, 4-21x faster than the reference at flattened cosine 0.9991-0.9999; isolation script; HTTP byte-identical to main with every new path off; make test-gemma2 (vision tower golden 0.999106 native, 0.999257 forced tier; audio 0.99992). M5 Max: media test 5-18x at 0.9992-0.9999, make test-gemma2 with the vision tower golden at 0.999007 (0.999154 with the scalar graph) against the 0.999 gate. EI_GEMMA2_SEMANTICS=official runs clean; it moves audio by 0.005 and video by 0.15 against the default, as the mode is meant to. Co-authored-by: Eric Hartford <eric.hartford@lazarus.enterprises>
…l prepack Ported from Yazeedx0/embeddinggemma-cpu-max (PR #6). On the CPU backend a media tower held inference_mutex for the whole tower (seconds for a 1120-token image) and every text wave queued behind it waited. Media requests now also take cpu_media_mutex ahead of inference_mutex, so one media request at a time is inside the towers, and the towers call the yield hook between layers: when text waves are waiting (text_engine_lock counts them) the tower releases inference_mutex until one of them has run or 50 ms pass, then takes its turn again. Text waves never take cpu_media_mutex or a tower lock, so the handoff cannot deadlock. Installed when the CPU backend serves media without GPU media slots; EI_GEMMA2_CPU_YIELD=0 keeps the plain lock. EI_GEMMA2_CPU_PREPACK=1 runs one text, image and audio input before the server listens so the one-time weight repacks are not paid by the first request of each modality (117 ms on a Ryzen 9950X). Off by default: the repacked copies of towers a deployment never uses would stay resident. Validated: make test-gemma2 (HTTP suite with concurrency and mixed batches) on an M5 Max and a Ryzen 9950X; HTTP responses byte-identical to main with the new paths off. Co-authored-by: Eric Hartford <eric.hartford@lazarus.enterprises>
The fast CPU media towers route their Q8 GEMMs through the repacked kernels, so on AVX-512 CPUs every vision and audio weight gained a second copy: server RSS after one text, image and audio request rose from main's 931 MB to 1536 MB on a Ryzen 9950X. The VNNI256 and sgemm repacks already drop their source pages (gemma2_release.h); the pre-existing 512-bit repack now does the same, and RSS is 925 MB (904 MB anonymous, 21 MB file). Numerics are untouched: HTTP responses stay byte-identical to main with the new paths off, and make test-gemma2 is unchanged.
CONTRIBUTING: the new EI_GEMMA2_CPU_* flags, the semantics switch, the kernel and media test targets, the per-tier test routine and the isolation rule. perf/optimization_status.md: what landed from PR #6, the 9950X and M5 Max numbers, the head-of-line and RSS probes, the fixes made while porting and what was left out. AGENTS.md: the lessons. CI: the linux-x86_64 job runs test_gemma2_attn and test_gemma2_vnni256 at the default, avx2 and scalar tiers and builds with clang; the ARM jobs run both tests on their native routes.
|
Note Currently processing new changes in this PR. This may take a few minutes, please wait... ⚙️ Run configuration
📒 Files selected for processing: 31
✨ Finishing Touches
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Re-applies the engine work from @Yazeedx0's PR #6 (fork
embeddinggemma-cpu-max) on main as a curated series of commits, with the contributor as author. No branding, nocpu-max/tree, no--versionor identity line; every "cpu-max fork" comment rewritten. Validated step by step on an M5 Max, a Ryzen 9950X (GCC 15) and an EPYC 9754 (GCC 11).What lands
EI_GEMMA2_CPU_FLASH(on)EI_GEMMA2_CPU_VNNI256,EI_GEMMA2_CPU_RELEASE_WEIGHTS(on)EI_GEMMA2_CPU_SGEMM(on)EI_GEMMA2_CPU_MEDIA_FAST(on)EI_GEMMA2_CPU_YIELD(on),EI_GEMMA2_CPU_PREPACK(off)Deliberate differences from PR #6
log(mel + 0.001), no video timestamps) is opt-in viaEI_GEMMA2_SEMANTICS=official; Metal and CUDA vision still read the GGUF flag, so a default flip is a separate PR.weight_cache_lockkeeps its arch guard (the fork broke Linux aarch64).aligned_allocwrapped for the UCRT. All new AVX functionsg2x_-prefixed socheck-avx2-isolation.pyholds. Audio q/k/v clamps per tensor instead ofei_die. Thread-local profiler accumulators.EI_GEMMA2_CPU_X86=avxvnniadded so AVX-512 hosts can exercise the 256-bit kernel.Numbers (9950X unless noted)
make test-gemma2: text golden minimum 0.999474 (main 0.999641), vision tower 0.999106 (0.999159), audio 0.99992. On the M5 Max the fast vision tower is at 0.999007 against the 0.999 gate (0.999154 with the scalar graph): the thinnest margin in the suite, flagged in the log.Closes #6 (superseded; the engine work is credited to its author in each commit).
Summary by CodeRabbit
New Features
Documentation