diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 94b8d60..6c609d6 100644 --- a/.github/workflows/ci.yml +++ b/.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} diff --git a/AGENTS.md b/AGENTS.md index 894fae0..bf5ddc7 100644 --- a/AGENTS.md +++ b/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 diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md index 6885bb0..c55777b 100644 --- a/CONTRIBUTING.md +++ b/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 diff --git a/Makefile b/Makefile index 8444182..3a6ac26 100644 --- a/Makefile +++ b/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 diff --git a/perf/optimization_status.md b/perf/optimization_status.md index 75c4668..46e8499 100644 --- a/perf/optimization_status.md +++ b/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`. + diff --git a/src/common.h b/src/common.h index 4a13fa0..65ebb27 100644 --- a/src/common.h +++ b/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 +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 */ diff --git a/src/engine.c b/src/engine.c index 34b54ab..9a742cd 100644 --- a/src/engine.c +++ b/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; diff --git a/src/gemma2.c b/src/gemma2.c index 0314c50..3c7d41b 100644 --- a/src/gemma2.c +++ b/src/gemma2.c @@ -1,6 +1,86 @@ +#if defined(__linux__) +#define _POSIX_C_SOURCE 200809L /* clock_gettime under -std=c11 */ +#endif #include "gemma2.h" #include "kernels.h" #include +#include + +/* EI_GEMMA2_REF_PROFILE=1: wall time per operation class of the reference + * graphs, printed per tower call. Off by default; no effect when off. The + * accumulators are per thread, so concurrent tower calls each report their own. */ +enum { RP_QUANT, RP_MM_Q8, RP_MM_F32, RP_ATTN, RP_NORM, RP_ROPE, RP_ACT, RP_ELEM, RP_CONV, RP_ALLOC, RP_N }; +static const char *rp_names[RP_N] = {"act_quant", "mm_q8", "mm_f32", "attention", "norm", "rope", + "activation", "elementwise", "conv", "alloc"}; +static _Thread_local double rp_ms[RP_N]; +static int rp_on(void) { + static int on = -1; + if (on < 0) { + const char *e = getenv("EI_GEMMA2_REF_PROFILE"); + on = e && strcmp(e, "0") != 0; + } + return on; +} +static double rp_now(void) { + struct timespec ts; + clock_gettime(CLOCK_MONOTONIC, &ts); + return ts.tv_sec * 1e3 + ts.tv_nsec * 1e-6; +} +#define RP(cat, ...) \ + do { \ + if (!rp_on()) { \ + __VA_ARGS__; \ + } else { \ + double rp_t0_ = rp_now(); \ + __VA_ARGS__; \ + rp_ms[cat] += rp_now() - rp_t0_; \ + } \ + } while (0) +static void rp_report(const char *tower, size_t tokens, double total) { + if (!rp_on()) + return; + double sum = 0; + for (int i = 0; i < RP_N; i++) + sum += rp_ms[i]; + fprintf(stderr, "ref-profile %s tokens=%zu total=%.1f ms:", tower, tokens, total); + for (int i = 0; i < RP_N; i++) + fprintf(stderr, " %s=%.1f(%.1f%%)", rp_names[i], rp_ms[i], 100 * rp_ms[i] / (total > 0 ? total : 1)); + fprintf(stderr, " other=%.1f\n", total - sum); + memset(rp_ms, 0, sizeof rp_ms); +} + +/* Model semantics (see gemma2.h). The default is llama.cpp's behaviour, which + * the goldens in testdata/ were produced with; EI_GEMMA2_SEMANTICS=official + * selects the official processor's. Read once. */ +bool ei_gemma2_official_semantics(void) { + static int on = -1; + if (on < 0) { + const char *e = getenv("EI_GEMMA2_SEMANTICS"); + on = e && !strcmp(e, "official"); + } + return on; +} +/* Vision MLP activation: EI_GEMMA2_VISION_GELU=tanh|quick overrides; else the + * official semantics use tanh GELU and the default follows the GGUF's + * clip.use_gelu, which this model's conversion leaves unset (QuickGELU). */ +bool ei_gemma2_vision_tanh_gelu(const ei_gguf *g) { + static int mode = -1; + if (mode < 0) { + const char *e = getenv("EI_GEMMA2_VISION_GELU"); + mode = !e ? 2 : !strcmp(e, "tanh") ? 1 : 0; + } + if (mode != 2) + return mode; + return ei_gemma2_official_semantics() || ei_gguf_kv_bool(g, "clip.use_gelu", false); +} +/* Installed once at startup by a server that shares the engine pool between + * CPU media towers and text waves; the towers call it between layers. */ +static void (*cpu_yield_hook)(void); +void ei_gemma2_set_cpu_yield(void (*hook)(void)) { cpu_yield_hook = hook; } +void ei_gemma2_cpu_yield(void) { + if (cpu_yield_hook) + cpu_yield_hook(); +} /* Tensor storage is [input, output]; activations are contiguous token rows. */ static const ei_tensor *tensor(const ei_gguf *g, const char *name) { @@ -144,14 +224,14 @@ static void mm(ei_thread_pool *pool, const ei_tensor *w, const float *x, int n, ei_block_q8_0 *q = NULL; if (w->type == EI_T_Q8_0) { q = ei_xcalloc((size_t)n * k / 32, sizeof(*q)); - for (int t = 0; t < n; ++t) - ei_quantize_row_q8_0(x + (size_t)t * k, q + (size_t)t * k / 32, k); + RP(RP_QUANT, for (int t = 0; t < n; ++t) + ei_quantize_row_q8_0(x + (size_t)t * k, q + (size_t)t * k / 32, k)); } mm_job job = {w, x, q, y, k, m}; - ei_parallel_for(pool, n * m, 32, mm_rows, &job); + RP(w->type == EI_T_Q8_0 ? RP_MM_Q8 : RP_MM_F32, ei_parallel_for(pool, n * m, 32, mm_rows, &job)); free(q); } -static void norm(float *x, int n, int d, const ei_tensor *w) { +static void norm_body(float *x, int n, int d, const ei_tensor *w) { for (int t = 0; t < n; ++t) { float *r = x + (size_t)t * d; double ss = 0; @@ -162,6 +242,7 @@ static void norm(float *x, int n, int d, const ei_tensor *w) { r[z] *= scale * (w ? value(w, z) : 1.0f); } } +static void norm(float *x, int n, int d, const ei_tensor *w) { RP(RP_NORM, norm_body(x, n, d, w)); } static void rope(float *x, int n, int heads, int d, float base) { for (int t = 0; t < n; ++t) for (int h = 0; h < heads; ++h) { @@ -216,7 +297,7 @@ static void attention_rows(void *opaque, int32_t begin, int32_t end) { static void attention(ei_thread_pool *pool, const float *q, const float *k, const float *v, int n, int heads, int kv, int d, int window, float *out) { attn_job j = {q, k, v, out, n, heads, kv, d, window}; - ei_parallel_for(pool, n * heads, 1, attention_rows, &j); + RP(RP_ATTN, ei_parallel_for(pool, n * heads, 1, attention_rows, &j)); } bool ei_gemma2_embed(const ei_gguf *g, ei_thread_pool *pool, const int32_t *ids, const float *inputs, size_t count, float out[768], char *err, size_t err_len) { @@ -302,21 +383,29 @@ bool ei_gemma2_embed(const ei_gguf *g, ei_thread_pool *pool, const int32_t *ids, /* Gemma4 vision tower: RGB patch projection, axial RoPE, 3x3 mean pool. */ bool ei_gemma2_vision(const ei_gguf *g, ei_thread_pool *pool, const uint8_t *rgb, int width, int height, float **output, size_t *tokens, char *err, size_t err_len) { + if (ei_gemma2_cpu_media_fast()) + return ei_gemma2_vision_fast(g, pool, rgb, width, height, output, tokens, err, err_len); + return ei_gemma2_vision_ref(g, pool, rgb, width, height, output, tokens, err, err_len); +} +/* The scalar reference graph (EI_GEMMA2_CPU_MEDIA_FAST=0). */ +bool ei_gemma2_vision_ref(const ei_gguf *g, ei_thread_pool *pool, const uint8_t *rgb, int width, + int height, float **output, size_t *tokens, char *err, size_t err_len) { if (width < 48 || height < 48 || width % 48 || height % 48 || (size_t)width * height > 1120 * 48 * 48) { snprintf(err, err_len, "vision input must be aligned to 48 pixels and at most 1120 soft tokens"); return false; } + double rp_total = rp_on() ? rp_now() : 0; int cols = width / 16, rows = height / 16, n = cols * rows; - bool use_gelu = ei_gguf_kv_bool(g, "clip.use_gelu", false); + bool use_gelu = ei_gemma2_vision_tanh_gelu(g); /* GGUF flag, or the semantics switch */ float *patches = alloc_f((size_t)n * 768), *x = alloc_f((size_t)n * 768), *h = alloc_f((size_t)n * 768); float *q = alloc_f((size_t)n * 768), *k = alloc_f((size_t)n * 768), *v = alloc_f((size_t)n * 768); float *a = alloc_f((size_t)n * 768), *up = alloc_f((size_t)n * 3072), *gate = alloc_f((size_t)n * 3072); - for (int t = 0; t < n; ++t) + RP(RP_ELEM, for (int t = 0; t < n; ++t) for (int c = 0; c < 3; ++c) for (int y = 0; y < 16; ++y) for (int z = 0; z < 16; ++z) @@ -324,23 +413,23 @@ bool ei_gemma2_vision(const ei_gguf *g, ei_thread_pool *pool, const uint8_t *rgb 2.0f * rgb[((size_t)(t / cols * 16 + y) * width + t % cols * 16 + z) * 3 + c] / 255.0f - - 1; + 1); ei_tensor patch = *tensor(g, "v.patch_embd.weight"); patch.ne[0] = 768; patch.ne[1] = 768; - for (size_t i = 0; i < (size_t)n * 768; ++i) - patches[i] = ei_fp16_to_fp32(ei_fp32_to_fp16(patches[i])); + RP(RP_ELEM, for (size_t i = 0; i < (size_t)n * 768; ++i) + patches[i] = ei_fp16_to_fp32(ei_fp32_to_fp16(patches[i]))); mm(pool, &patch, patches, n, x); free(patches); const ei_tensor *pos = tensor(g, "v.position_embd.weight"); - for (int t = 0; t < n; ++t) + RP(RP_ELEM, for (int t = 0; t < n; ++t) for (int z = 0; z < 768; ++z) x[(size_t)t * 768 + z] += value(pos, (size_t)(t % cols) * 768 + z) + - value(pos, ((size_t)10240 + t / cols) * 768 + z); + value(pos, ((size_t)10240 + t / cols) * 768 + z)); for (int l = 0; l < 16; ++l) { char name[128]; #define V(suffix) (snprintf(name, sizeof name, "v.blk.%d." suffix ".weight", l), tensor(g, name)) - memcpy(h, x, (size_t)n * 768 * sizeof(float)); + RP(RP_ELEM, memcpy(h, x, (size_t)n * 768 * sizeof(float))); norm(h, n, 768, V("ln1")); mm(pool, V("attn_q"), h, n, q); mm(pool, V("attn_k"), h, n, k); @@ -348,7 +437,7 @@ bool ei_gemma2_vision(const ei_gguf *g, ei_thread_pool *pool, const uint8_t *rgb norm(q, n * 12, 64, V("attn_q_norm")); norm(k, n * 12, 64, V("attn_k_norm")); norm(v, n * 12, 64, NULL); - for (int t = 0; t < n; ++t) + RP(RP_ROPE, for (int t = 0; t < n; ++t) for (int head = 0; head < 12; ++head) for (int axis = 0; axis < 2; ++axis) for (int z = 0; z < 16; ++z) { @@ -361,36 +450,36 @@ bool ei_gemma2_vision(const ei_gguf *g, ei_thread_pool *pool, const uint8_t *rgb q[i + 16] = qa * si + qb * co; k[i] = ka * co - kb * si; k[i + 16] = ka * si + kb * co; - } + }); attention(pool, q, k, v, n, 12, 12, 64, 0, a); mm(pool, V("attn_out"), a, n, h); norm(h, n, 768, V("attn_post_norm")); - for (size_t i = 0; i < (size_t)n * 768; ++i) - x[i] += h[i]; - memcpy(h, x, (size_t)n * 768 * sizeof(float)); + RP(RP_ELEM, for (size_t i = 0; i < (size_t)n * 768; ++i) + x[i] += h[i]); + RP(RP_ELEM, memcpy(h, x, (size_t)n * 768 * sizeof(float))); norm(h, n, 768, V("ln2")); mm(pool, V("ffn_up"), h, n, up); mm(pool, V("ffn_gate"), h, n, gate); - for (size_t i = 0; i < (size_t)n * 3072; ++i) + RP(RP_ACT, for (size_t i = 0; i < (size_t)n * 3072; ++i) gate[i] = (use_gelu ? ei_gelu_tanh(gate[i]) : gate[i] / (1.0f + expf(-1.702f * gate[i]))) * - up[i]; + up[i]); mm(pool, V("ffn_down"), gate, n, h); norm(h, n, 768, V("ffn_post_norm")); - for (size_t i = 0; i < (size_t)n * 768; ++i) - x[i] += h[i]; + RP(RP_ELEM, for (size_t i = 0; i < (size_t)n * 768; ++i) + x[i] += h[i]); #undef V } int out_cols = cols / 3, out_rows = rows / 3, nt = out_cols * out_rows; memset(h, 0, (size_t)nt * 768 * sizeof(float)); - for (int t = 0; t < nt; ++t) + RP(RP_ELEM, for (int t = 0; t < nt; ++t) for (int yy = 0; yy < 3; ++yy) for (int xx = 0; xx < 3; ++xx) for (int z = 0; z < 768; ++z) h[(size_t)t * 768 + z] += x[((size_t)(t / out_cols * 3 + yy) * cols + t % out_cols * 3 + xx) * 768 + z] / - 9.0f * sqrtf(768.0f); + 9.0f * sqrtf(768.0f)); norm(h, nt, 768, NULL); *output = alloc_f((size_t)nt * 512); *tokens = nt; @@ -403,6 +492,7 @@ bool ei_gemma2_vision(const ei_gguf *g, ei_thread_pool *pool, const uint8_t *rgb free(a); free(up); free(gate); + rp_report("vision", (size_t)n, rp_on() ? rp_now() - rp_total : 0); return true; } @@ -421,25 +511,33 @@ static void clipped_mm(const ei_gguf *g, ei_thread_pool *pool, const char *name, const ei_tensor *w = tensor(g, name); size_t count = (size_t)n * w->ne[0]; float *clamped = alloc_f(count); - for (size_t i = 0; i < count; ++i) - clamped[i] = fminf(hi, fmaxf(lo, x[i])); + RP(RP_ELEM, for (size_t i = 0; i < count; ++i) + clamped[i] = fminf(hi, fmaxf(lo, x[i]))); mm(pool, w, clamped, n, y); free(clamped); snprintf(key, sizeof key, "%.*s.output_min", (int)len, name); lo = scalar_or(g, key, -INFINITY); snprintf(key, sizeof key, "%.*s.output_max", (int)len, name); hi = scalar_or(g, key, INFINITY); - for (size_t i = 0; i < (size_t)n * w->ne[1]; ++i) - y[i] = fminf(hi, fmaxf(lo, y[i])); + RP(RP_ELEM, for (size_t i = 0; i < (size_t)n * w->ne[1]; ++i) + y[i] = fminf(hi, fmaxf(lo, y[i]))); } static float silu(float x) { return x / (1.0f + expf(-x)); } /* Audio features are time-major, 128 HTK log-magnitude mel bins. */ bool ei_gemma2_audio(const ei_gguf *g, ei_thread_pool *pool, const float *mel, int frames, float **output, size_t *tokens, char *err, size_t err_len) { + if (ei_gemma2_cpu_media_fast()) + return ei_gemma2_audio_fast(g, pool, mel, frames, output, tokens, err, err_len); + return ei_gemma2_audio_ref(g, pool, mel, frames, output, tokens, err, err_len); +} +/* The scalar reference graph (EI_GEMMA2_CPU_MEDIA_FAST=0). */ +bool ei_gemma2_audio_ref(const ei_gguf *g, ei_thread_pool *pool, const float *mel, int frames, + float **output, size_t *tokens, char *err, size_t err_len) { if (frames < 1 || frames > 3000) { snprintf(err, err_len, "audio chunk must contain 1..3000 mel frames"); return false; } + double rp_total = rp_on() ? rp_now() : 0; int nt = frames, nf = 128, nc = 1; float *cur = alloc_f((size_t)nt * nf); memcpy(cur, mel, (size_t)nt * nf * sizeof(float)); @@ -449,7 +547,7 @@ bool ei_gemma2_audio(const ei_gguf *g, ei_thread_pool *pool, const float *mel, i const ei_tensor *w = tensor(g, name); int ot = (nt + 1) / 2, of = (nf + 1) / 2, oc = (int)w->ne[3]; float *next = alloc_f((size_t)ot * of * oc); - for (int t = 0; t < ot; ++t) + RP(RP_CONV, for (int t = 0; t < ot; ++t) for (int f = 0; f < of; ++f) for (int o = 0; o < oc; ++o) { float sum = 0; @@ -463,10 +561,10 @@ bool ei_gemma2_audio(const ei_gguf *g, ei_thread_pool *pool, const float *mel, i value(w, ((size_t)o * nc + c) * 9 + ky * 3 + kx); } next[((size_t)t * of + f) * oc + o] = sum; - } + }); snprintf(name, sizeof name, "a.conv1d.%d.norm.weight", layer); const ei_tensor *weight = tensor(g, name); - for (int t = 0; t < ot * of; ++t) { + RP(RP_NORM, for (int t = 0; t < ot * of; ++t) { float *r = next + (size_t)t * oc; double mean = 0, var = 0; for (int c = 0; c < oc; ++c) @@ -477,7 +575,7 @@ bool ei_gemma2_audio(const ei_gguf *g, ei_thread_pool *pool, const float *mel, i float scale = 1.0f / sqrtf((float)(var / oc) + 1e-6f); for (int c = 0; c < oc; ++c) r[c] = fmaxf(0, ((float)(r[c] - mean)) * scale * value(weight, c)); - } + }); free(cur); cur = next; nt = ot; @@ -505,28 +603,28 @@ bool ei_gemma2_audio(const ei_gguf *g, ei_thread_pool *pool, const float *mel, i snprintf(name, sizeof name, "a.blk.%d." suffix ".weight", l); \ clipped_mm(g, pool, name, input, n, output); \ } while (0) - memcpy(h, x, count * sizeof(float)); + RP(RP_ELEM, memcpy(h, x, count * sizeof(float))); norm(h, n, 1024, A("ffn_norm")); AM("ffn_up", h, up); - for (size_t i = 0; i < (size_t)n * 4096; ++i) - up[i] = silu(up[i]); + RP(RP_ACT, for (size_t i = 0; i < (size_t)n * 4096; ++i) + up[i] = silu(up[i])); AM("ffn_down", up, h); norm(h, n, 1024, A("ffn_post_norm")); - for (size_t i = 0; i < count; ++i) - x[i] += 0.5f * h[i]; - memcpy(h, x, count * sizeof(float)); + RP(RP_ELEM, for (size_t i = 0; i < count; ++i) + x[i] += 0.5f * h[i]); + RP(RP_ELEM, memcpy(h, x, count * sizeof(float))); norm(h, n, 1024, A("attn_pre_norm")); AM("attn_q", h, q); AM("attn_k", h, k); AM("attn_v", h, v); const ei_tensor *qs = A("per_dim_scale"); - for (size_t i = 0; i < count; ++i) { + RP(RP_ELEM, for (size_t i = 0; i < count; ++i) { q[i] *= value(qs, i % 128) / (sqrtf(128.0f) * logf(2.0f)); k[i] *= log1pf(expf(1.0f)) / logf(2.0f); - } + }); mm(pool, A("attn_k_rel"), pos, 13, rpe); - memset(a, 0, count * sizeof(float)); - for (int t = 0; t < n; ++t) + RP(RP_ELEM, memset(a, 0, count * sizeof(float))); + RP(RP_ATTN, for (int t = 0; t < n; ++t) for (int head = 0; head < 8; ++head) { float scores[12], mx = -INFINITY; int lo = t > 11 ? t - 11 : 0; @@ -550,52 +648,52 @@ bool ei_gemma2_audio(const ei_gguf *g, ei_thread_pool *pool, const float *mel, i for (int z = 0; z < 128; ++z) a[(size_t)t * 1024 + head * 128 + z] += scores[s - lo] / total * v[(size_t)s * 1024 + head * 128 + z]; - } + }); AM("attn_out", a, h); norm(h, n, 1024, A("attn_post_norm")); - for (size_t i = 0; i < count; ++i) - x[i] += h[i]; + RP(RP_ELEM, for (size_t i = 0; i < count; ++i) + x[i] += h[i]); /* GGUF conversion swaps norm_conv and conv_norm; match the reference loader. */ - memcpy(h, x, count * sizeof(float)); + RP(RP_ELEM, memcpy(h, x, count * sizeof(float))); norm(h, n, 1024, A("conv_norm")); AM("conv_pw1", h, up); - for (int t = 0; t < n; ++t) + RP(RP_ACT, for (int t = 0; t < n; ++t) for (int z = 0; z < 1024; ++z) h[(size_t)t * 1024 + z] = - up[(size_t)t * 2048 + z] / (1 + expf(-up[(size_t)t * 2048 + 1024 + z])); + up[(size_t)t * 2048 + z] / (1 + expf(-up[(size_t)t * 2048 + 1024 + z]))); const ei_tensor *dw = A("conv_dw"); - for (int t = 0; t < n; ++t) + RP(RP_CONV, for (int t = 0; t < n; ++t) for (int z = 0; z < 1024; ++z) { float sum = 0; for (int j = 0; j < 5; ++j) if (t + j >= 4) sum += h[(size_t)(t + j - 4) * 1024 + z] * value(dw, z * 5 + j); a[(size_t)t * 1024 + z] = sum; - } + }); norm(a, n, 1024, A("norm_conv")); - for (size_t i = 0; i < count; ++i) - a[i] = silu(a[i]); + RP(RP_ACT, for (size_t i = 0; i < count; ++i) + a[i] = silu(a[i])); AM("conv_pw2", a, h); - for (size_t i = 0; i < count; ++i) - x[i] += h[i]; - memcpy(h, x, count * sizeof(float)); + RP(RP_ELEM, for (size_t i = 0; i < count; ++i) + x[i] += h[i]); + RP(RP_ELEM, memcpy(h, x, count * sizeof(float))); norm(h, n, 1024, A("ffn_norm_1")); AM("ffn_up_1", h, up); - for (size_t i = 0; i < (size_t)n * 4096; ++i) - up[i] = silu(up[i]); + RP(RP_ACT, for (size_t i = 0; i < (size_t)n * 4096; ++i) + up[i] = silu(up[i])); AM("ffn_down_1", up, h); norm(h, n, 1024, A("ffn_post_norm_1")); - for (size_t i = 0; i < count; ++i) - x[i] += 0.5f * h[i]; + RP(RP_ELEM, for (size_t i = 0; i < count; ++i) + x[i] += 0.5f * h[i]); norm(x, n, 1024, A("ln2")); #undef AM #undef A } clipped_mm(g, pool, "a.pre_encode.out.weight", x, n, up); const ei_tensor *bias = tensor(g, "a.pre_encode.out.bias"); - for (int t = 0; t < n; ++t) + RP(RP_ELEM, for (int t = 0; t < n; ++t) for (int z = 0; z < 1536; ++z) - up[(size_t)t * 1536 + z] += value(bias, z); + up[(size_t)t * 1536 + z] += value(bias, z)); norm(up, n, 1536, NULL); *output = alloc_f((size_t)n * 512); *tokens = n; @@ -609,5 +707,6 @@ bool ei_gemma2_audio(const ei_gguf *g, ei_thread_pool *pool, const float *mel, i free(up); free(pos); free(rpe); + rp_report("audio", (size_t)frames, rp_on() ? rp_now() - rp_total : 0); return true; } diff --git a/src/gemma2.h b/src/gemma2.h index bfd742a..eb47cc0 100644 --- a/src/gemma2.h +++ b/src/gemma2.h @@ -23,8 +23,42 @@ bool ei_gemma2_embed_batch_fast_ctx(ei_gemma2_fast *ctx, const ei_gguf *g, ei_th const int32_t *ids, const float *inputs, const size_t *offsets, size_t batch, float *out, char *err, size_t err_len); +/* Fast-graph kernels reused by the CPU media towers. The GEMM + * uses ctx's activation-repack workspace; callers serialize per context. */ +void ei_gemma2_fast_gemm_q8(ei_gemma2_fast *ctx, ei_thread_pool *pool, const ei_tensor *w, + const ei_block_q8_0 *xq, int n, float *y); +void ei_gemma2_fast_gemm_f(ei_thread_pool *pool, const ei_tensor *w, const float *x, int n, float *y); +void ei_gemma2_fast_quant_row(const float *x, ei_block_q8_0 *q, int width); +float ei_gemma2_fast_sum_squares(const float *x, int d); +void ei_gemma2_fast_gelu_mul(float *g, const float *u, int n); bool ei_gemma2_vision(const ei_gguf *g, ei_thread_pool *pool, const uint8_t *rgb, int width, int height, float **output, size_t *tokens, char *err, size_t err_len); +/* SIMD/threaded CPU vision tower (default; EI_GEMMA2_CPU_MEDIA_FAST=0 + * restores the reference graph above). Same math as the reference: Q8 activations + * for Q8 weights, fp32 attention/norms, the selected GELU variant. */ +bool ei_gemma2_cpu_media_fast(void); +/* Model semantics. The default matches llama.cpp (QuickGELU in the vision MLP, + * log(max(mel, floor)), "MM:SS " text before video frames); EI_GEMMA2_SEMANTICS=official + * selects the official processor's behaviour (tanh GELU, log(mel + floor), no + * timestamps). EI_GEMMA2_VISION_GELU, EI_GEMMA2_AUDIO_LOGMEL and + * EI_GEMMA2_VIDEO_TIMESTAMPS override one item each. */ +bool ei_gemma2_official_semantics(void); +/* The CPU media towers call ei_gemma2_cpu_yield() between + * layers; a server that shares the engine pool with text waves installs a + * hook there that lets waiting waves run (EI_GEMMA2_CPU_YIELD=0: never). */ +void ei_gemma2_set_cpu_yield(void (*hook)(void)); +void ei_gemma2_cpu_yield(void); +bool ei_gemma2_vision_tanh_gelu(const ei_gguf *g); +bool ei_gemma2_vision_ref(const ei_gguf *g, ei_thread_pool *pool, const uint8_t *rgb, int width, + int height, float **output, size_t *tokens, char *err, size_t err_len); +bool ei_gemma2_vision_fast(const ei_gguf *g, ei_thread_pool *pool, const uint8_t *rgb, int width, + int height, float **output, size_t *tokens, char *err, size_t err_len); bool ei_gemma2_audio(const ei_gguf *g, ei_thread_pool *pool, const float *mel, int frames, float **output, size_t *tokens, char *err, size_t err_len); +/* SIMD/threaded CPU audio tower (default; EI_GEMMA2_CPU_MEDIA_FAST=0 + * restores the reference graph). Same math as the reference. */ +bool ei_gemma2_audio_ref(const ei_gguf *g, ei_thread_pool *pool, const float *mel, int frames, + float **output, size_t *tokens, char *err, size_t err_len); +bool ei_gemma2_audio_fast(const ei_gguf *g, ei_thread_pool *pool, const float *mel, int frames, + float **output, size_t *tokens, char *err, size_t err_len); #endif diff --git a/src/gemma2_audio_fast.c b/src/gemma2_audio_fast.c new file mode 100644 index 0000000..8e2ca13 --- /dev/null +++ b/src/gemma2_audio_fast.c @@ -0,0 +1,577 @@ +/* SIMD + threaded CPU audio tower for EmbeddingGemma 2. + * + * The reference CPU audio tower is the scalar graph (gemma2.c: + * per-element GEMM loops, both subsampling convolutions through a per-element + * type switch, single-threaded norms/clamps/activations/quantization). This + * path keeps the reference math -- Q8_0 activations for Q8_0 weights, the + * clipped-linear input/output clamps, fp32 convolutions, norms and local + * attention, the reference's tensor mapping -- and executes it with: + * - the subsampling convolutions as channel-vectorized AVX2 kernels with + * fused LayerNorm + ReLU, + * - Q8_0 GEMMs on the fast graph's routes (AVX-VNNI 256 repacked here), + * - per-token fused row passes (norm + clamp + quantize, clamp + SiLU + + * quantize, GLU, causal depthwise conv + norm + SiLU + quantize, + * post-norm + residual), threaded over tokens, + * - relative-position keys computed once per layer and cached, + * - a persistent workspace. + * EI_GEMMA2_CPU_MEDIA_FAST=0 selects the reference graph. */ +#include "gemma2.h" +#include "gemma2_fast_x86.h" +#include "gemma2_media_x86.h" +#include "kernels.h" +#include +#include +#include +#include +#if defined(__x86_64__) || defined(_M_X64) +#include +#define AVX2 __attribute__((target("avx2,fma,f16c"))) +#define UNROLL _Pragma("GCC unroll 8") +#endif + +enum { AD = 1024, AF = 4096, AHEADS = 8, AHD = 128, CTX = 12, NREL = 13 }; + +/* EI_GEMMA2_CPU_PROFILE=1: per-stage wall time of each audio call on stderr. */ +#include +enum { AP_CONV, AP_GEMM, AP_ROWS, AP_ATTN, AP_N }; +static const char *ap_names[AP_N] = {"subsample", "gemm", "rows", "attention"}; +static double ap_ms[AP_N]; +static int ap_on(void) { + static int on = -1; + if (on < 0) { + const char *e = getenv("EI_GEMMA2_CPU_PROFILE"); + on = e && strcmp(e, "0") != 0; + } + return on; +} +static double ap_now(void) { + struct timespec ts; + clock_gettime(CLOCK_MONOTONIC, &ts); + return ts.tv_sec * 1e3 + ts.tv_nsec * 1e-6; +} +#define AP(i, ...) \ + do { \ + double ap_t0_ = ap_on() ? ap_now() : 0; \ + __VA_ARGS__; \ + if (ap_on()) \ + ap_ms[i] += ap_now() - ap_t0_; \ + } while (0) + +static const ei_tensor *tensor_of(const ei_gguf *g, const char *name) { + return ei_gguf_tensor(g, name, true); +} +static const float *f32_of(const ei_gguf *g, const char *name) { + const ei_tensor *t = tensor_of(g, name); + if (t->type != EI_T_F32) + ei_die("EmbeddingGemma 2 audio: %s must be F32", name); + return t->data; +} +static float scalar_or(const ei_gguf *g, const char *name, float fallback) { + const ei_tensor *t = ei_gguf_tensor(g, name, false); + return t ? ((const float *)t->data)[0] : fallback; +} + +/* A clipped linear: Q8_0 weight plus optional input/output clamps. */ +typedef struct { + const ei_tensor *w; + float in_lo, in_hi, out_lo, out_hi; +} clinear; +static clinear clinear_of(const ei_gguf *g, const char *base) { + char key[160]; + clinear c; + snprintf(key, sizeof key, "%s.weight", base); + c.w = tensor_of(g, key); + snprintf(key, sizeof key, "%s.input_min", base); + c.in_lo = scalar_or(g, key, -INFINITY); + snprintf(key, sizeof key, "%s.input_max", base); + c.in_hi = scalar_or(g, key, INFINITY); + snprintf(key, sizeof key, "%s.output_min", base); + c.out_lo = scalar_or(g, key, -INFINITY); + snprintf(key, sizeof key, "%s.output_max", base); + c.out_hi = scalar_or(g, key, INFINITY); + return c; +} + +/* ---- workspace ------------------------------------------------------------ */ + +typedef struct { + size_t cap; /* tokens */ + float *x, *h, *hn, *q, *k, *v, *a; /* [cap][1024] */ + float *up; /* [cap][4096] */ + ei_block_q8_0 *xq; /* [cap][4096/32] */ + size_t conv1_cap; + float *c1; /* subsampled activations [t1][32][32] */ + float rel[12][NREL * AD]; /* per-layer relative keys */ + const void *rel_key; /* the projector's conv0 weight data: stable per mapping */ + float qscale[12][AHD], kscale; + float w0t[9][128]; /* conv0 weights, [tap][out] */ + float *w1t; /* conv1 weights, [tap][in][out] = [9][128][32] */ +} audio_ws; + +static pthread_mutex_t audio_lock = PTHREAD_MUTEX_INITIALIZER; +static audio_ws aws; +static ei_gemma2_fast *gemm_ctx; + +static void reserve(audio_ws *w, size_t n, size_t conv1_floats) { + if (w->cap < n) { + free(w->x), free(w->h), free(w->hn), free(w->q), free(w->k), free(w->v), free(w->a); + free(w->up), free(w->xq); + size_t row = (size_t)AD * sizeof(float); + w->x = ei_xmalloc(n * row); + w->h = ei_xmalloc(n * row); + w->hn = ei_xmalloc(n * row); + w->q = ei_xmalloc(n * row); + w->k = ei_xmalloc(n * row); + w->v = ei_xmalloc(n * row); + w->a = ei_xmalloc(n * row); + w->up = ei_xmalloc(n * AF * sizeof(float)); + w->xq = ei_xmalloc(n * (AF / 32) * sizeof(ei_block_q8_0)); + w->cap = n; + } + if (w->conv1_cap < conv1_floats) { + free(w->c1); + w->c1 = ei_xmalloc(conv1_floats * sizeof(float)); + w->conv1_cap = conv1_floats; + } +} + +/* Constant per projector: relative keys of the 13 positions (the reference + * recomputes them per call with the same Q8 GEMM), q/k scales, conv weights. */ +static void prepare(audio_ws *w, const ei_gguf *g, ei_thread_pool *pool) { + const void *key = tensor_of(g, "a.conv1d.0.weight")->data; + if (w->rel_key == key) + return; + float pos[NREL * AD]; + for (int p = 0; p < NREL; ++p) + for (int z = 0; z < 512; ++z) { + float theta = (12 - p) * expf(-z * logf(10000.0f) / 511); + pos[p * AD + z] = sinf(theta); + pos[p * AD + z + 512] = cosf(theta); + } + ei_block_q8_0 posq[NREL * AD / 32]; + for (int p = 0; p < NREL; ++p) + ei_quantize_row_q8_0(pos + p * AD, posq + p * (AD / 32), AD); + for (int l = 0; l < 12; ++l) { + char name[128]; + snprintf(name, sizeof name, "a.blk.%d.attn_k_rel.weight", l); + ei_gemma2_fast_gemm_q8(gemm_ctx, pool, tensor_of(g, name), posq, NREL, w->rel[l]); + snprintf(name, sizeof name, "a.blk.%d.per_dim_scale.weight", l); + const float *qs = f32_of(g, name); + for (int z = 0; z < AHD; ++z) + w->qscale[l][z] = qs[z] / (sqrtf(128.0f) * logf(2.0f)); + } + w->kscale = log1pf(expf(1.0f)) / logf(2.0f); + const float *c0 = f32_of(g, "a.conv1d.0.weight"), *c1 = f32_of(g, "a.conv1d.1.weight"); + for (int o = 0; o < 128; ++o) + for (int tap = 0; tap < 9; ++tap) + w->w0t[tap][o] = c0[(size_t)o * 9 + tap]; + free(w->w1t); + w->w1t = ei_xmalloc((size_t)9 * 128 * 32 * sizeof(float)); + for (int o = 0; o < 32; ++o) + for (int c = 0; c < 128; ++c) + for (int tap = 0; tap < 9; ++tap) + w->w1t[((size_t)tap * 128 + c) * 32 + o] = c1[((size_t)o * 128 + c) * 9 + tap]; + w->rel_key = key; +} + +/* ---- subsampling convolutions ------------------------------------------------ */ + + +static void layer_norm_relu(float *r, int c, const float *w) { + double mean = 0, var = 0; + for (int i = 0; i < c; ++i) + mean += r[i]; + mean /= c; + for (int i = 0; i < c; ++i) + var += (r[i] - mean) * (r[i] - mean); + float scale = 1.0f / sqrtf((float)(var / c) + 1e-6f); + for (int i = 0; i < c; ++i) + r[i] = fmaxf(0, ((float)(r[i] - mean)) * scale * w[i]); +} + +/* conv0: 1 -> 128 channels, 3x3, stride 2, pad 1, over (time, mel): one + * output time row (64 mel positions x 128 channels), LayerNorm + ReLU. */ +static void conv0_row(const audio_ws *w, const float *mel, int frames, int t, const float *norm, float *dst) { + for (int f = 0; f < 64; f++) { + float acc[128] = {0}; + for (int ky = 0; ky < 3; ky++) + for (int kx = 0; kx < 3; kx++) { + int ti = t * 2 + ky - 1, fi = f * 2 + kx - 1; + if (ti < 0 || ti >= frames || fi < 0 || fi >= 128) + continue; + float x = mel[(size_t)ti * 128 + fi]; + const float *wr = w->w0t[ky * 3 + kx]; + for (int o = 0; o < 128; o++) + acc[o] += x * wr[o]; + } + float *d = dst + (size_t)f * 128; + memcpy(d, acc, sizeof acc); + layer_norm_relu(d, 128, norm); + } +} + +#if defined(__x86_64__) || defined(_M_X64) +/* conv1 output (t, f) = 32 channels in 4 ymm; 1152 taps x inputs. `in` holds + * conv0 rows [row0, ...); bounds use the global time index. */ +/* g2x_ prefix: scripts/check-avx2-isolation.py's dispatched set. */ +AVX2 static void g2x_conv1_point(const float *in, int row0, int nt, int nf, int t, int f, const float *w1t, + float *out) { + __m256 acc[4]; + UNROLL for (int o = 0; o < 4; o++) acc[o] = _mm256_setzero_ps(); + for (int ky = 0; ky < 3; ky++) + for (int kx = 0; kx < 3; kx++) { + int ti = t * 2 + ky - 1, fi = f * 2 + kx - 1; + if (ti < 0 || ti >= nt || fi < 0 || fi >= nf) + continue; + const float *x = in + ((size_t)(ti - row0) * nf + fi) * 128; + const float *wr = w1t + (size_t)(ky * 3 + kx) * 128 * 32; + for (int c = 0; c < 128; c++) { + __m256 xb = _mm256_broadcast_ss(x + c); + UNROLL for (int o = 0; o < 4; o++) + acc[o] = _mm256_fmadd_ps(xb, _mm256_loadu_ps(wr + c * 32 + o * 8), acc[o]); + } + } + UNROLL for (int o = 0; o < 4; o++) _mm256_storeu_ps(out + o * 8, acc[o]); +} +#endif + +/* Both subsampling convolutions, streamed: a task producing conv1 rows + * [begin, end) computes the conv0 rows they read (2*begin-1 .. 2*end-1) into a + * task-local buffer instead of materializing conv0 for the whole clip (49 MB + * at 30 s); boundary rows are recomputed by neighbouring tasks. */ +typedef struct { + const audio_ws *w; + const float *mel, *norm0, *norm1; + float *out; + int frames, t0, t1; +} conv_job; +static void conv01_rows(void *opaque, int32_t begin, int32_t end) { + conv_job *j = opaque; + int r0 = 2 * begin - 1 < 0 ? 0 : 2 * begin - 1, r1 = 2 * (end - 1) + 2 < j->t0 ? 2 * (end - 1) + 2 : j->t0; + static _Thread_local float *buf; /* per pool thread, kept: <= (2 * grain + 1) conv0 rows */ + static _Thread_local size_t cap; + size_t need = (size_t)(r1 - r0) * 64 * 128; + if (cap < need) { + free(buf); + buf = ei_xmalloc(need * sizeof(float)); + cap = need; + } + for (int r = r0; r < r1; r++) + conv0_row(j->w, j->mel, j->frames, r, j->norm0, buf + (size_t)(r - r0) * 64 * 128); + for (int t = begin; t < end; t++) + for (int f = 0; f < 32; f++) { + float *dst = j->out + ((size_t)t * 32 + f) * 32; +#if defined(__x86_64__) || defined(_M_X64) + if (g2x_level() >= G2X_AVX2) + g2x_conv1_point(buf, r0, j->t0, 64, t, f, j->w->w1t, dst); + else +#endif + { + float acc[32] = {0}; + for (int tap = 0; tap < 9; tap++) { + int ti = t * 2 + tap / 3 - 1, fi = f * 2 + tap % 3 - 1; + if (ti < 0 || ti >= j->t0 || fi < 0 || fi >= 64) + continue; + const float *x = buf + ((size_t)(ti - r0) * 64 + fi) * 128; + for (int c = 0; c < 128; c++) + for (int o = 0; o < 32; o++) + acc[o] += x[c] * j->w->w1t[((size_t)tap * 128 + c) * 32 + o]; + } + memcpy(dst, acc, sizeof acc); + } + layer_norm_relu(dst, 32, j->norm1); + } +} + +/* ---- row passes ------------------------------------------------------------- */ + +typedef enum { + R_NORM, /* hn = rms(x) * w */ + R_QUANT, /* xq = q8(clamp(src, lo, hi)), width wd */ + R_SILU_QUANT, /* up = silu(clamp(up, olo, ohi)); xq = q8(clamp(up, lo, hi)) */ + R_POST_ADD, /* x += sc * rms(clamp(h, olo, ohi)) * w */ + R_ADD, /* x += clamp(h, olo, ohi) (light-conv output, no norm) */ + R_QKV, /* clamp q/k/v outputs (own bounds each); scale q per dim, k by kscale */ + R_GLU, /* h = a * sigmoid(b), a/b = clamp(up)[0:1024] / [1024:2048] */ + R_DWCONV, /* a = rms(causal dwconv5(h)) * w; a = silu(a); xq = q8(clamp(a)) */ + R_FINAL_NORM, /* x = rms(x) * w */ + R_BIAS_NORM_QUANT, /* up = rms(clamp(up) + bias) (1536); xq = q8(clamp(up, lo, hi)) */ +} arow; +typedef struct { + arow op; + audio_ws *w; + const float *nw, *bias; + float *src; + int wd, layer; + float lo, hi, olo, ohi, sc; + float klo, khi, vlo, vhi; /* R_QKV: k and v output clamps (q uses olo/ohi) */ +} arow_job; + +static void rms_into(float *y, const float *x, int d, const float *w) { + float scale = 1.0f / sqrtf(ei_gemma2_fast_sum_squares(x, d) / d + 1e-6f); + if (w) + for (int z = 0; z < d; z++) + y[z] = x[z] * (scale * w[z]); + else + for (int z = 0; z < d; z++) + y[z] = x[z] * scale; +} +static void clamp_row(float *x, int d, float lo, float hi) { + if (lo == -INFINITY && hi == INFINITY) + return; + for (int z = 0; z < d; z++) + x[z] = fminf(hi, fmaxf(lo, x[z])); +} +static void quant_clamped(const float *x, int d, float lo, float hi, ei_block_q8_0 *q) { + if (lo == -INFINITY && hi == INFINITY) { + ei_gemma2_fast_quant_row(x, q, d); + return; + } + float tmp[AF]; + for (int z = 0; z < d; z++) + tmp[z] = fminf(hi, fmaxf(lo, x[z])); + ei_gemma2_fast_quant_row(tmp, q, d); +} +static void arow_tokens(void *opaque, int32_t begin, int32_t end) { + arow_job *j = opaque; + audio_ws *w = j->w; + float tmp[AF]; + for (int t = begin; t < end; t++) { + float *x = w->x + (size_t)t * AD; + switch (j->op) { + case R_NORM: + rms_into(w->hn + (size_t)t * AD, x, AD, j->nw); + break; + case R_QUANT: + quant_clamped(j->src + (size_t)t * j->wd, j->wd, j->lo, j->hi, w->xq + (size_t)t * (j->wd / 32)); + break; + case R_SILU_QUANT: { + float *u = w->up + (size_t)t * AF; + clamp_row(u, AF, j->olo, j->ohi); + g2m_silu(u, u, AF); + quant_clamped(u, AF, j->lo, j->hi, w->xq + (size_t)t * (AF / 32)); + break; + } + case R_POST_ADD: { + float *h = w->h + (size_t)t * AD; + clamp_row(h, AD, j->olo, j->ohi); + rms_into(tmp, h, AD, j->nw); + for (int z = 0; z < AD; z++) + x[z] += j->sc * tmp[z]; + break; + } + case R_ADD: { + float *h = w->h + (size_t)t * AD; + clamp_row(h, AD, j->olo, j->ohi); + for (int z = 0; z < AD; z++) + x[z] += h[z]; + break; + } + case R_QKV: { + float *q = w->q + (size_t)t * AD, *k = w->k + (size_t)t * AD, *v = w->v + (size_t)t * AD; + clamp_row(q, AD, j->olo, j->ohi); + clamp_row(k, AD, j->klo, j->khi); + clamp_row(v, AD, j->vlo, j->vhi); + const float *qs = w->qscale[j->layer]; + for (int z = 0; z < AD; z++) { + q[z] *= qs[z % AHD]; + k[z] *= w->kscale; + } + break; + } + case R_GLU: { + float *u = w->up + (size_t)t * 2 * AD, *h = w->h + (size_t)t * AD; /* pw1 rows are 2048 wide */ + clamp_row(u, 2 * AD, j->olo, j->ohi); + for (int z = 0; z < AD; z++) + h[z] = u[z] / (1 + expf(-u[AD + z])); + break; + } + case R_DWCONV: { + const float *dw = j->bias; /* [1024][5] */ + float *a = w->a + (size_t)t * AD; + for (int z = 0; z < AD; z++) { + float sum = 0; + for (int i = 0; i < 5; i++) + if (t + i >= 4) + sum += w->h[(size_t)(t + i - 4) * AD + z] * dw[z * 5 + i]; + a[z] = sum; + } + rms_into(a, a, AD, j->nw); + g2m_silu(a, a, AD); + quant_clamped(a, AD, j->lo, j->hi, w->xq + (size_t)t * (AD / 32)); + break; + } + case R_FINAL_NORM: + rms_into(x, x, AD, j->nw); + break; + case R_BIAS_NORM_QUANT: { + float *u = w->up + (size_t)t * 1536; + clamp_row(u, 1536, j->olo, j->ohi); + for (int z = 0; z < 1536; z++) + u[z] += j->bias[z]; + rms_into(u, u, 1536, NULL); + quant_clamped(u, 1536, j->lo, j->hi, w->xq + (size_t)t * (1536 / 32)); + break; + } + } + } +} +static void arows(ei_thread_pool *pool, int n, arow_job job) { ei_parallel_for(pool, n, 4, arow_tokens, &job); } + +/* ---- local relative-position attention ---------------------------------------- */ + +typedef struct { + audio_ws *w; + int n, layer; +} aattn_job; +static void aattn_tokens(void *opaque, int32_t begin, int32_t end) { + aattn_job *j = opaque; + audio_ws *w = j->w; + const float *rel = w->rel[j->layer]; + for (int t = begin; t < end; t++) + for (int hd = 0; hd < AHEADS; hd++) { + const float *q = w->q + (size_t)t * AD + hd * AHD; + float scores[CTX], mx = -INFINITY; + int lo = t > CTX - 1 ? t - (CTX - 1) : 0; + for (int s = lo; s <= t; s++) { + const float *k = w->k + (size_t)s * AD + hd * AHD, *r = rel + (size_t)(12 - t + s) * AD + hd * AHD; + float dot = 0; + for (int z = 0; z < AHD; z++) + dot += q[z] * (k[z] + r[z]); + dot = 50 * tanhf(dot / 50); + scores[s - lo] = dot; + mx = dot > mx ? dot : mx; + } + float total = 0; + for (int s = lo; s <= t; s++) { + scores[s - lo] = expf(scores[s - lo] - mx); + total += scores[s - lo]; + } + float *o = w->a + (size_t)t * AD + hd * AHD; + memset(o, 0, AHD * sizeof(float)); + for (int s = lo; s <= t; s++) { + float p = scores[s - lo] / total; + const float *v = w->v + (size_t)s * AD + hd * AHD; + for (int z = 0; z < AHD; z++) + o[z] += p * v[z]; + } + } +} + +/* ---- tower --------------------------------------------------------------------- */ + +static void create_ctx(void) { gemm_ctx = ei_gemma2_fast_create(); } + +static void gemm(ei_thread_pool *pool, const clinear *c, const ei_block_q8_0 *xq, int n, float *y) { + AP(AP_GEMM, ei_gemma2_fast_gemm_q8(gemm_ctx, pool, c->w, xq, n, y)); +} + +bool ei_gemma2_audio_fast(const ei_gguf *g, ei_thread_pool *pool, const float *mel, int frames, + float **output, size_t *tokens, char *err, size_t err_len) { + if (frames < 1 || frames > 3000) { + snprintf(err, err_len, "audio chunk must contain 1..3000 mel frames"); + return false; + } + static pthread_once_t once = PTHREAD_ONCE_INIT; + pthread_once(&once, create_ctx); + pthread_mutex_lock(&audio_lock); + ei_cpu_binding binding; /* hybrid pools: the caller joins on a P-core */ + ei_thread_pool_bind_caller(pool, &binding); + audio_ws *w = &aws; + const int t0 = (frames + 1) / 2, t1 = (t0 + 1) / 2, n = t1; /* conv0: t0 x 64 mel x 128 ch; conv1: t1 x 32 x 32 */ + reserve(w, (size_t)n, (size_t)t1 * 32 * 32); + prepare(w, g, pool); + + /* Subsampling: (frames x 128) -> (t0 x 64 x 128) -> (t1 x 32 x 32), LayerNorm + ReLU. */ + conv_job cj = {w, mel, f32_of(g, "a.conv1d.0.norm.weight"), f32_of(g, "a.conv1d.1.norm.weight"), w->c1, + frames, t0, t1}; + AP(AP_CONV, ei_parallel_for(pool, t1, 4, conv01_rows, &cj)); + AP(AP_GEMM, ei_gemma2_fast_gemm_f(pool, tensor_of(g, "a.input_projection.weight"), w->c1, n, w->x)); + + for (int l = 0; l < 12; ++l) { + char base[96], name[128]; + if (l) + ei_gemma2_cpu_yield(); /* between layers: waiting text waves may run */ +#define CL(suffix) (snprintf(base, sizeof base, "a.blk.%d." suffix, l), clinear_of(g, base)) +#define AF32(suffix) (snprintf(name, sizeof name, "a.blk.%d." suffix ".weight", l), f32_of(g, name)) + clinear up = CL("ffn_up"), down = CL("ffn_down"), q = CL("attn_q"), k = CL("attn_k"), v = CL("attn_v"); + clinear out = CL("attn_out"), pw1 = CL("conv_pw1"), pw2 = CL("conv_pw2"); + clinear up1 = CL("ffn_up_1"), down1 = CL("ffn_down_1"); + /* Half-step FFN 1. */ + AP(AP_ROWS, arows(pool, n, (arow_job){.op = R_NORM, .w = w, .nw = AF32("ffn_norm")}); + arows(pool, n, (arow_job){.op = R_QUANT, .w = w, .src = w->hn, .wd = AD, .lo = up.in_lo, .hi = up.in_hi})); + gemm(pool, &up, w->xq, n, w->up); + AP(AP_ROWS, arows(pool, n, (arow_job){.op = R_SILU_QUANT, .w = w, .olo = up.out_lo, .ohi = up.out_hi, + .lo = down.in_lo, .hi = down.in_hi})); + gemm(pool, &down, w->xq, n, w->h); + AP(AP_ROWS, arows(pool, n, (arow_job){.op = R_POST_ADD, .w = w, .nw = AF32("ffn_post_norm"), .sc = 0.5f, + .olo = down.out_lo, .ohi = down.out_hi})); + /* Local attention with relative positions. */ + AP(AP_ROWS, arows(pool, n, (arow_job){.op = R_NORM, .w = w, .nw = AF32("attn_pre_norm")})); + const clinear *qkv[3] = {&q, &k, &v}; + float *dst[3] = {w->q, w->k, w->v}; + float qlo = NAN, qhi = NAN; + for (int i = 0; i < 3; i++) { + if (!(qkv[i]->in_lo == qlo && qkv[i]->in_hi == qhi)) /* q/k/v usually share clamps */ + AP(AP_ROWS, arows(pool, n, (arow_job){.op = R_QUANT, .w = w, .src = w->hn, .wd = AD, + .lo = qkv[i]->in_lo, .hi = qkv[i]->in_hi})); + qlo = qkv[i]->in_lo, qhi = qkv[i]->in_hi; + gemm(pool, qkv[i], w->xq, n, dst[i]); + } + AP(AP_ROWS, arows(pool, n, (arow_job){.op = R_QKV, .w = w, .layer = l, .olo = q.out_lo, .ohi = q.out_hi, + .klo = k.out_lo, .khi = k.out_hi, .vlo = v.out_lo, .vhi = v.out_hi})); + aattn_job aj = {w, n, l}; + AP(AP_ATTN, ei_parallel_for(pool, n, 4, aattn_tokens, &aj)); + AP(AP_ROWS, arows(pool, n, (arow_job){.op = R_QUANT, .w = w, .src = w->a, .wd = AD, .lo = out.in_lo, + .hi = out.in_hi})); + gemm(pool, &out, w->xq, n, w->h); + AP(AP_ROWS, arows(pool, n, (arow_job){.op = R_POST_ADD, .w = w, .nw = AF32("attn_post_norm"), .sc = 1.0f, + .olo = out.out_lo, .ohi = out.out_hi})); + /* Light convolution. The GGUF conversion swaps norm_conv and conv_norm; + * the reference loader's mapping is used (conv_norm before pw1). */ + AP(AP_ROWS, arows(pool, n, (arow_job){.op = R_NORM, .w = w, .nw = AF32("conv_norm")}); + arows(pool, n, (arow_job){.op = R_QUANT, .w = w, .src = w->hn, .wd = AD, .lo = pw1.in_lo, .hi = pw1.in_hi})); + gemm(pool, &pw1, w->xq, n, w->up); + AP(AP_ROWS, arows(pool, n, (arow_job){.op = R_GLU, .w = w, .olo = pw1.out_lo, .ohi = pw1.out_hi}); + arows(pool, n, (arow_job){.op = R_DWCONV, .w = w, .bias = AF32("conv_dw"), .nw = AF32("norm_conv"), + .lo = pw2.in_lo, .hi = pw2.in_hi})); + gemm(pool, &pw2, w->xq, n, w->h); + AP(AP_ROWS, arows(pool, n, (arow_job){.op = R_ADD, .w = w, .olo = pw2.out_lo, .ohi = pw2.out_hi})); + /* Half-step FFN 2, then the block's output norm. */ + AP(AP_ROWS, arows(pool, n, (arow_job){.op = R_NORM, .w = w, .nw = AF32("ffn_norm_1")}); + arows(pool, n, (arow_job){.op = R_QUANT, .w = w, .src = w->hn, .wd = AD, .lo = up1.in_lo, .hi = up1.in_hi})); + gemm(pool, &up1, w->xq, n, w->up); + AP(AP_ROWS, arows(pool, n, (arow_job){.op = R_SILU_QUANT, .w = w, .olo = up1.out_lo, .ohi = up1.out_hi, + .lo = down1.in_lo, .hi = down1.in_hi})); + gemm(pool, &down1, w->xq, n, w->h); + AP(AP_ROWS, arows(pool, n, (arow_job){.op = R_POST_ADD, .w = w, .nw = AF32("ffn_post_norm_1"), .sc = 0.5f, + .olo = down1.out_lo, .ohi = down1.out_hi}); + arows(pool, n, (arow_job){.op = R_FINAL_NORM, .w = w, .nw = AF32("ln2")})); +#undef CL +#undef AF32 + } + clinear pre = clinear_of(g, "a.pre_encode.out"), proj = clinear_of(g, "mm.a.input_projection"); + AP(AP_ROWS, arows(pool, n, (arow_job){.op = R_QUANT, .w = w, .src = w->x, .wd = AD, .lo = pre.in_lo, + .hi = pre.in_hi})); + gemm(pool, &pre, w->xq, n, w->up); + AP(AP_ROWS, arows(pool, n, (arow_job){.op = R_BIAS_NORM_QUANT, .w = w, .bias = f32_of(g, "a.pre_encode.out.bias"), + .olo = pre.out_lo, .ohi = pre.out_hi, .lo = proj.in_lo, .hi = proj.in_hi})); + *output = ei_xmalloc((size_t)n * 512 * sizeof(float)); + *tokens = (size_t)n; + gemm(pool, &proj, w->xq, n, *output); + for (size_t i = 0; i < (size_t)n * 512; i++) + (*output)[i] = fminf(proj.out_hi, fmaxf(proj.out_lo, (*output)[i])); + if (ap_on()) { + double total = 0; + for (int i = 0; i < AP_N; i++) + total += ap_ms[i]; + fprintf(stderr, "audio-fast frames=%d tokens=%d total=%.1f ms:", frames, n, total); + for (int i = 0; i < AP_N; i++) + fprintf(stderr, " %s=%.1f", ap_names[i], ap_ms[i]); + fprintf(stderr, "\n"); + memset(ap_ms, 0, sizeof ap_ms); + } + ei_thread_pool_unbind_caller(&binding); + pthread_mutex_unlock(&audio_lock); + return true; +} diff --git a/src/gemma2_fast.c b/src/gemma2_fast.c index db4163f..6d600d0 100644 --- a/src/gemma2_fast.c +++ b/src/gemma2_fast.c @@ -9,6 +9,9 @@ * tile, threaded SIMD row kernels, cached RoPE tables, persistent workspace. */ #include "gemma2.h" #include "gemma2_fast_x86.h" +#include "gemma2_media_x86.h" +#include "gemma2_vnni256.h" +#include "gemma2_sgemm_x86.h" #include "kernels.h" #include #include @@ -241,12 +244,16 @@ static bool gemm_q8_i8mm(workspace *ws, ei_thread_pool *pool, const ei_tensor *w float *y); static bool gemm_q8_x86_512(workspace *ws, ei_thread_pool *pool, const ei_tensor *w, const ei_block_q8_0 *xq, int n, float *y); +static bool gemm_q8_x86_vnni256(workspace *ws, ei_thread_pool *pool, const ei_tensor *w, + const ei_block_q8_0 *xq, int n, float *y); static void gemm_q8(workspace *ws, ei_thread_pool *pool, const ei_tensor *w, const ei_block_q8_0 *xq, int n, float *y) { if (gemm_q8_i8mm(ws, pool, w, xq, n, y)) return; if (gemm_q8_x86_512(ws, pool, w, xq, n, y)) return; + if (gemm_q8_x86_vnni256(ws, pool, w, xq, n, y)) + return; // Short waves: 8-row blocks so every thread gets work; long: 32-row tiles. int rr = n <= 8 ? 8 : 32; q8_job j = {w->data, xq, y, (int)w->ne[0], (int)w->ne[1], n, (int)(w->ne[1] + rr - 1) / rr, rr}; @@ -474,6 +481,82 @@ static bool gemm_q8_x86_512(workspace *ws, ei_thread_pool *pool, const ei_tensor } #endif +/* ---- x86 AVX-VNNI 256-bit GEMM on 8-row repacked weights ---------------- + * gemma2_vnni256.c. Tasks are 32 rows (4 groups) x 32 tokens; inside a task the + * group's weight panel (8 x k bytes) stays in L1 across the token loop. */ +#if defined(__x86_64__) +typedef struct { + const void *w; + const ei_block_q8_0 *xq; + uint8_t *xu; + float *xs; + float *y; + int k, m, n, row_tasks; +} v256_job; +static void v256_acts(void *opaque, int32_t begin, int32_t end) { + v256_job *j = opaque; + g2x_acts(j->xq, begin, end, j->k, j->xu, j->xs); +} +/* Short waves (<= 8 tokens): one task per 8-row group. */ +static void v256_groups(void *opaque, int32_t begin, int32_t end) { + v256_job *j = opaque; + for (int g = begin; g < end; g++) + for (int t = 0; t < j->n; t += G2V_TILE_TOKENS) + g2v_tile(j->w, g, j->xu, j->xs, t, j->n - t < G2V_TILE_TOKENS ? j->n - t : G2V_TILE_TOKENS, j->k, + j->y, j->m); +} +static void v256_rows(void *opaque, int32_t begin, int32_t end) { + v256_job *j = opaque; + for (int task = begin; task < end; task++) { + int rb = task % j->row_tasks, tb = task / j->row_tasks; + int g_end = rb * 4 + 4 < j->m / 8 ? rb * 4 + 4 : j->m / 8; + int t_end = tb * 32 + 32 < j->n ? tb * 32 + 32 : j->n; + for (int g = rb * 4; g < g_end; g++) + for (int t = tb * 32; t < t_end; t += G2V_TILE_TOKENS) + g2v_tile(j->w, g, j->xu, j->xs, t, t_end - t < G2V_TILE_TOKENS ? t_end - t : G2V_TILE_TOKENS, + j->k, j->y, j->m); + } +} +static bool gemm_q8_x86_vnni256(workspace *ws, ei_thread_pool *pool, const ei_tensor *w, + const ei_block_q8_0 *xq, int n, float *y) { + int k = (int)w->ne[0], m = (int)w->ne[1]; + // EI_GEMMA2_CPU_VNNI256_SHORT=0 keeps waves of <= 8 tokens on the generic tile. + static int short_waves = -1; + if (short_waves < 0) { + const char *e = getenv("EI_GEMMA2_CPU_VNNI256_SHORT"); + short_waves = !(e && !strcmp(e, "0")); + } + if (!g2v_enabled() || (n <= 8 && !short_waves) || k % 32 || m % 8) + return false; + pthread_mutex_lock(&weight_cache_lock); + const void *pw = g2v_weight_of(w->data, w->data, m, k, pool); + pthread_mutex_unlock(&weight_cache_lock); + if (!pw) + return false; + if ((size_t)n * k > ws->x512_cap) { + free(ws->x512_u), free(ws->x512_s); + ws->x512_u = ei_xmalloc((size_t)n * k); + ws->x512_s = ei_xmalloc((size_t)n * (k / 32) * sizeof(float)); + ws->x512_cap = (size_t)n * k; + } + v256_job j = {pw, xq, ws->x512_u, ws->x512_s, y, k, m, n, (m / 8 + 3) / 4}; + if (n <= 8) { + g2x_acts(xq, 0, n, k, ws->x512_u, ws->x512_s); + PROF(0, ei_parallel_for(pool, m / 8, short_grain(pool, m / 8, n), v256_groups, &j)); + return true; + } + PROF(0, ei_parallel_for(pool, n, 16, v256_acts, &j)); + PROF(0, ei_parallel_for(pool, j.row_tasks * ((n + 31) / 32), 1, v256_rows, &j)); + return true; +} +#else +static bool gemm_q8_x86_vnni256(workspace *ws, ei_thread_pool *pool, const ei_tensor *w, + const ei_block_q8_0 *xq, int n, float *y) { + (void)ws, (void)pool, (void)w, (void)xq, (void)n, (void)y; + return false; +} +#endif + /* Float-activation GEMM for BF16/F16/F32 weights (the PLE projection). */ typedef struct { const ei_tensor *w; @@ -559,7 +642,48 @@ static void f_rows(void *opaque, int32_t begin, int32_t end) { } } } +/* x86 packed fp32 GEMM (gemma2_sgemm_x86.c): tasks are 4 panels (64 outputs) + * x 48 tokens. EI_GEMMA2_CPU_SGEMM=0 restores f_rows. The repack cache is + * keyed by the weight data (stable per mapping; callers pass tensor copies + * with adjusted shapes, so the struct address is not a key). */ +#if defined(__x86_64__) +typedef struct { + const void *w; + const float *x; + float *y; + int n, panel_tasks, panels; +} sgemm_job; +static void sgemm_tasks(void *opaque, int32_t begin, int32_t end) { + sgemm_job *j = opaque; + for (int task = begin; task < end; task++) { + int pb = task % j->panel_tasks, tb = task / j->panel_tasks; + int p0 = pb * 4, np = j->panels - p0 < 4 ? j->panels - p0 : 4; + int t0 = tb * 48, t1 = t0 + 48 < j->n ? t0 + 48 : j->n; + g2s_block(j->w, j->x, t0, t1, p0, np, j->y); + } +} +static bool gemm_f_sgemm(ei_thread_pool *pool, const ei_tensor *w, const float *x, int n, float *y) { + int k = (int)w->ne[0], m = (int)w->ne[1]; + if (!g2s_enabled() || m % 16) + return false; + pthread_mutex_lock(&weight_cache_lock); + const void *pw = g2s_weight_of(w->data, w->data, (int)w->type, m, k, pool); + pthread_mutex_unlock(&weight_cache_lock); + if (!pw) + return false; + sgemm_job j = {pw, x, y, n, (m / 16 + 3) / 4, m / 16}; + PROF(1, ei_parallel_for(pool, j.panel_tasks * ((n + 47) / 48), 1, sgemm_tasks, &j)); + return true; +} +#else +static bool gemm_f_sgemm(ei_thread_pool *pool, const ei_tensor *w, const float *x, int n, float *y) { + (void)pool, (void)w, (void)x, (void)n, (void)y; + return false; +} +#endif static void gemm_f(ei_thread_pool *pool, const ei_tensor *w, const float *x, int n, float *y) { + if (gemm_f_sgemm(pool, w, x, n, y)) + return; f_job j = {w, x, y, (int)w->ne[0], (int)w->ne[1], n}; if (j.k > 1024 || j.k % 4) ei_die("gemm_f supports k <= 1024, multiple of 4"); @@ -895,6 +1019,80 @@ static void attention_rows(void *opaque, int32_t begin, int32_t end) { } } +/* ---- shared with the CPU media towers ---------------------------------- */ + +void ei_gemma2_fast_gemm_q8(ei_gemma2_fast *ctx, ei_thread_pool *pool, const ei_tensor *w, + const ei_block_q8_0 *xq, int n, float *y) { + gemm_q8(&ctx->ws, pool, w, xq, n, y); +} +void ei_gemma2_fast_gemm_f(ei_thread_pool *pool, const ei_tensor *w, const float *x, int n, float *y) { + gemm_f(pool, w, x, n, y); +} +void ei_gemma2_fast_quant_row(const float *x, ei_block_q8_0 *q, int width) { quant_row(x, q, width); } +float ei_gemma2_fast_sum_squares(const float *x, int d) { return sum_squares(x, d); } +void ei_gemma2_fast_gelu_mul(float *g, const float *u, int n) { gelu_mul_row(g, u, n); } + +/* ---- x86 tiled online-softmax attention ------------------------------------ + * gemma2_media_x86.c's g2m_attn over (head, query block) tasks: per-query key + * bounds carry the packed-sequence limits and the +-512 sliding window, so + * one kernel serves every layer. x86 AVX2 only (the NEON path above stays on + * ARM); EI_GEMMA2_CPU_FLASH=0 restores the attention above. */ +#if defined(__x86_64__) +static bool flash_enabled(void) { + static int cached = -1; + if (cached < 0) { + const char *flag = getenv("EI_GEMMA2_CPU_FLASH"); + cached = (!flag || strcmp(flag, "0")) && g2x_level() >= G2X_AVX2; + } + return cached; +} +typedef struct { + int n, kv, d, window, blocks; + workspace *ws; +} flash_job; +static void flash_tasks(void *opaque, int32_t begin, int32_t end) { + flash_job *j = opaque; + workspace *ws = j->ws; + const int qb = g2m_attn_qblock(j->d); + int32_t lo[G2M_ATTN_QB_MAX], hi[G2M_ATTN_QB_MAX]; + for (int task = begin; task < end; task++) { + int hd = task / j->blocks, q0 = task % j->blocks * qb, kh = hd / (4 / j->kv); + int nq = j->n - q0 < qb ? j->n - q0 : qb, klo = 1 << 30, khi = 0; + for (int i = 0; i < nq; i++) { + int t = q0 + i; + lo[i] = (int32_t)ws->lo[t], hi[i] = (int32_t)ws->hi[t]; + if (j->window) { + if (t - j->window > lo[i]) + lo[i] = t - j->window; + if (t + j->window + 1 < hi[i]) + hi[i] = t + j->window + 1; + } + klo = lo[i] < klo ? lo[i] : klo; + khi = hi[i] > khi ? hi[i] : khi; + } + g2m_attn_args a = {.q = ws->q, .k = ws->k, .v = ws->v, .out = ws->a, .qstride = (size_t)4 * j->d, + .kvstride = (size_t)j->kv * j->d, .ostride = (size_t)4 * j->d, + .qcol = (size_t)hd * j->d, .kvcol = (size_t)kh * j->d, .ocol = (size_t)hd * j->d, + .d = j->d, .q0 = q0, .nq = nq, .lo = lo, .hi = hi, .klo = klo, .khi = khi, + .scale = 1.0f}; + g2m_attn(&a); + } +} +static bool attention_flash(workspace *ws, ei_thread_pool *pool, int n, int kv, int d, int window) { + if (!flash_enabled()) + return false; + int qb = g2m_attn_qblock(d); + flash_job fj = {n, kv, d, window, (n + qb - 1) / qb, ws}; + PROF(3, ei_parallel_for(pool, 4 * fj.blocks, 1, flash_tasks, &fj)); + return true; +} +#else +static bool attention_flash(workspace *ws, ei_thread_pool *pool, int n, int kv, int d, int window) { + (void)ws, (void)pool, (void)n, (void)kv, (void)d, (void)window; + return false; +} +#endif + /* ---- graph -------------------------------------------------------------- */ ei_gemma2_fast *ei_gemma2_fast_create(void) { @@ -932,6 +1130,21 @@ bool ei_gemma2_embed_batch_fast_ctx(ei_gemma2_fast *ctx, const ei_gguf *g, ei_th workspace *ws = &ctx->ws; reserve(ws, count); int n = (int)count; + /* Hybrid CPUs: waves below EI_GEMMA2_CPU_HYBRID_TOKENS run on + * the pool's P-core threads only (every parallel step otherwise waits on + * the slower E-cores); larger waves use every core. The caller is pinned + * to the pool's first (P-core) CPU for the wave. */ + ei_cpu_binding binding; + ei_thread_pool_bind_caller(pool, &binding); + const int perf = ei_thread_pool_performance_threads(pool); + static int hybrid_tokens = -1; + if (hybrid_tokens < 0) { + const char *e = getenv("EI_GEMMA2_CPU_HYBRID_TOKENS"); + hybrid_tokens = e ? atoi(e) : 512; + } + const bool capped = perf > 0 && n < hybrid_tokens; + if (capped) + ei_thread_pool_set_cap(pool, perf); for (size_t s = 0; s < batch; s++) for (size_t t = offsets[s]; t < offsets[s + 1]; t++) { ws->lo[t] = (uint32_t)offsets[s]; @@ -970,7 +1183,8 @@ bool ei_gemma2_embed_batch_fast_ctx(ei_gemma2_fast *ctx, const ei_gguf *g, ei_th rows(ws, pool, n, (row_job){.op = ROW_QKV_POST, .d = d, .kv = kv, .rope_base = l % 6 == 5 ? 1 : 0}); attn_job aj = {n, kv, d, l % 6 == 5 ? 0 : 512, ws}; - PROF(3, ei_parallel_for(pool, 4 * ((n + QB - 1) / QB), 1, attention_rows, &aj)); + if (!attention_flash(ws, pool, n, kv, d, aj.window)) + PROF(3, ei_parallel_for(pool, 4 * ((n + QB - 1) / QB), 1, attention_rows, &aj)); rows(ws, pool, n, (row_job){.op = ROW_QUANT, .in = ws->a, .width = 4 * d}); gemm_q8(ws, pool, T("attn_output"), ws->xq, n, ws->h); dequant_norm(T("post_attention_norm"), 512, ws->norm_w[2]); @@ -1014,6 +1228,9 @@ bool ei_gemma2_embed_batch_fast_ctx(ei_gemma2_fast *ctx, const ei_gguf *g, ei_th fprintf(stderr, "\n"); } memset(prof, 0, sizeof prof); + if (capped) + ei_thread_pool_set_cap(pool, 0); + ei_thread_pool_unbind_caller(&binding); pthread_mutex_unlock(&ctx->lock); return true; } diff --git a/src/gemma2_fast_x86.c b/src/gemma2_fast_x86.c index 01389f8..3859e87 100644 --- a/src/gemma2_fast_x86.c +++ b/src/gemma2_fast_x86.c @@ -6,11 +6,16 @@ * (signed); negating INT8_MIN would wrap, but activations come from the * symmetric +-127 quantizer, so -128 never occurs in x. AVX-512 is promoted * per kernel by measurement, never blanket (perf/optimization_status.md). - * EI_GEMMA2_CPU_X86=scalar|avx2|vnni256 forces a tier (default: the best; - * on AVX-512 VNNI CPUs the repacked 512-bit GEMM). */ + * EI_GEMMA2_CPU_X86=scalar|avx2|vnni256|avxvnni forces a tier (default: the + * best; on AVX-512 VNNI CPUs the repacked 512-bit GEMM). */ +#if defined(__linux__) +#define _DEFAULT_SOURCE /* madvise under -std=c11 (gemma2_release.h) */ +#endif #include "gemma2_fast_x86.h" +#include "gemma2_release.h" #if defined(__x86_64__) || defined(_M_X64) +#include #include #include #include @@ -21,6 +26,16 @@ #define AVXVNNI __attribute__((target("avx2,fma,f16c,avxvnni"))) static int wide512; /* repacked 512-bit VNNI GEMM (g2x_gemm512_tile) */ +/* AVX-VNNI (VEX-encoded VNNI): CPUID.(EAX=7,ECX=1):EAX[4]. Read directly because + * __builtin_cpu_supports("avxvnni") is rejected by clang before 19. */ +static int cpu_avxvnni(void) { + unsigned a, b, c, d; + if (__get_cpuid_max(0, NULL) < 7) + return 0; + __cpuid_count(7, 1, a, b, c, d); + (void)b, (void)c, (void)d; + return (a >> 4) & 1; +} int g2x_level(void) { static int level = -1; if (level >= 0) @@ -30,7 +45,7 @@ int g2x_level(void) { __builtin_cpu_supports("f16c"); int vnni512 = avx2 && __builtin_cpu_supports("avx512vl") && __builtin_cpu_supports("avx512vnni"); - int avxvnni = avx2 && __builtin_cpu_supports("avxvnni"); + int avxvnni = avx2 && cpu_avxvnni(); /* avx2 implies OS-enabled YMM state */ level = vnni512 ? G2X_VNNI512 : avxvnni ? G2X_AVXVNNI : avx2 ? G2X_AVX2 : G2X_NONE; wide512 = vnni512 && __builtin_cpu_supports("avx512dq"); const char *force = getenv("EI_GEMMA2_CPU_X86"); @@ -40,6 +55,8 @@ int g2x_level(void) { level = G2X_AVX2, wide512 = 0; else if (force && !strcmp(force, "vnni256")) wide512 = 0; /* VNNI tier, 256-bit Q8 tile */ + else if (force && !strcmp(force, "avxvnni") && avxvnni) + level = G2X_AVXVNNI, wide512 = 0; /* the AVX-VNNI tier (repacked 256-bit GEMM) on an AVX-512 CPU */ return level; } int g2x_gemm512_enabled(void) { @@ -153,6 +170,8 @@ const void *g2x_weight_of(const void *key, const ei_block_q8_0 *blocks, int m, i sum += blocks[b].qs[z]; w->c[b] = -128 * sum; } + /* The repack is the only copy the kernel reads from now on. */ + g2_release_source(blocks, (size_t)m * kb * sizeof(ei_block_q8_0)); w->key = key; g2x_weight_count++; return w; diff --git a/src/gemma2_media_x86.c b/src/gemma2_media_x86.c new file mode 100644 index 0000000..57f24d4 --- /dev/null +++ b/src/gemma2_media_x86.c @@ -0,0 +1,295 @@ +/* SIMD kernels for the EmbeddingGemma 2 CPU text graph and media towers. + * + * Attention (g2m_attn64) is a tiled online-softmax ("flash") kernel written in + * a transposed form so that every softmax step is a vertical vector op: + * Q^T [64][QB] packed once per task (QB queries, zero padded to 16) + * S^T [KB][QB] = K_blk . Q^T, K rows read in place (broadcast per key) + * P^T = exp(S^T - m), running max m and sum l kept per query lane + * O^T [64][QB] += V_blk^T . P^T, V rows read in place (broadcast per dim) + * so neither K nor V is ever copied or transposed; the only packed operand is + * the task's Q block (16 KB). Arithmetic is fp32 throughout (no reduced + * precision); only summation order differs from the scalar reference. */ +#include "gemma2_media_x86.h" +#include "gemma2_fast_x86.h" + +#include +#include +#include + +/* ---- portable fallbacks -------------------------------------------------- */ + +int g2m_attn_qblock(int d) { + /* EI_G2M_QB=16..128 (multiple of 16) overrides the query block for widths <= 64, and up + * to 64 for widths <= 128 (tuning experiments). */ + static int qb = -1; + if (qb < 0) { + const char *e = getenv("EI_G2M_QB"); + int v = e ? atoi(e) : 0; + qb = v >= 16 && v <= 128 && v % 16 == 0 ? v : 0; + } + if (qb && (d <= 64 || (d <= 128 && qb <= 64))) + return qb; /* Q^T/O^T hold d x QB <= 8192 floats */ + return d <= 128 ? 64 : d <= 256 ? 32 : 16; +} + +void g2m_attn_reference(const g2m_attn_args *a) { + float scores[G2M_ATTN_KB], o[512]; + for (int i = 0; i < a->nq; i++) { + const float *qi = a->q + (size_t)(a->q0 + i) * a->qstride + a->qcol; + int lo = a->lo ? a->lo[i] : a->klo, hi = a->hi ? a->hi[i] : a->khi; + float m = -INFINITY, l = 0; + memset(o, 0, (size_t)a->d * sizeof(float)); + for (int b = lo; b < hi; b += G2M_ATTN_KB) { + int kn = hi - b < G2M_ATTN_KB ? hi - b : G2M_ATTN_KB; + float mb = -INFINITY; + for (int j = 0; j < kn; j++) { + const float *kj = a->k + (size_t)(b + j) * a->kvstride + a->kvcol; + float s = 0; + for (int dd = 0; dd < a->d; dd++) + s += qi[dd] * a->scale * kj[dd]; + scores[j] = s; + mb = s > mb ? s : mb; + } + float mn = mb > m ? mb : m, alpha = expf(m - mn); + l *= alpha; + for (int dd = 0; dd < a->d; dd++) + o[dd] *= alpha; + for (int j = 0; j < kn; j++) { + float p = expf(scores[j] - mn); + l += p; + const float *vj = a->v + (size_t)(b + j) * a->kvstride + a->kvcol; + for (int dd = 0; dd < a->d; dd++) + o[dd] += p * vj[dd]; + } + m = mn; + } + float *oi = a->out + (size_t)(a->q0 + i) * a->ostride + a->ocol; + for (int dd = 0; dd < a->d; dd++) + oi[dd] = l > 0 ? o[dd] / l : 0.0f; + } +} + +#if defined(__x86_64__) || defined(_M_X64) +#include +#define AVX2 __attribute__((target("avx2,fma,f16c"))) +/* Fixed-size register tiles must unroll at the upstream -O2 (GCC keeps such + * loops rolled there and spills the accumulator arrays). */ +#define UNROLL _Pragma("GCC unroll 8") + +/* exp via x = n ln2 + r, |r| <= ln2/2, degree-7 Taylor in Horner form. + * Inputs below -87.3 (including -inf, e.g. masked attention scores) and NaN + * return exactly 0; inputs above 88.7 saturate. Every AVX body here carries + * the g2x_ prefix: scripts/check-avx2-isolation.py allows AVX instructions + * only in the runtime-dispatched kernels. */ +AVX2 static inline __m256 g2x_exp8(__m256 x) { + const __m256 live = _mm256_cmp_ps(x, _mm256_set1_ps(-87.3f), _CMP_GE_OQ); + x = _mm256_min_ps(_mm256_max_ps(x, _mm256_set1_ps(-87.3f)), _mm256_set1_ps(88.7f)); + __m256 n = _mm256_round_ps(_mm256_mul_ps(x, _mm256_set1_ps(1.44269504088896341f)), + _MM_FROUND_TO_NEAREST_INT | _MM_FROUND_NO_EXC); + __m256 r = _mm256_fnmadd_ps(n, _mm256_set1_ps(0.693359375f), x); + r = _mm256_fnmadd_ps(n, _mm256_set1_ps(-2.12194440e-4f), r); + __m256 p = _mm256_set1_ps(1.0f / 5040); + p = _mm256_fmadd_ps(p, r, _mm256_set1_ps(1.0f / 720)); + p = _mm256_fmadd_ps(p, r, _mm256_set1_ps(1.0f / 120)); + p = _mm256_fmadd_ps(p, r, _mm256_set1_ps(1.0f / 24)); + p = _mm256_fmadd_ps(p, r, _mm256_set1_ps(1.0f / 6)); + p = _mm256_fmadd_ps(p, r, _mm256_set1_ps(0.5f)); + p = _mm256_fmadd_ps(p, r, _mm256_set1_ps(1.0f)); + p = _mm256_fmadd_ps(p, r, _mm256_set1_ps(1.0f)); + __m256i e = _mm256_slli_epi32(_mm256_add_epi32(_mm256_cvtps_epi32(n), _mm256_set1_epi32(127)), 23); + return _mm256_and_ps(_mm256_mul_ps(p, _mm256_castsi256_ps(e)), live); +} + +/* Scratch per thread: Q^T and O^T hold d x QB floats (<= 8192 each by the + * choice of QB), S^T/P^T KB x QB. 80 KB, so not on small worker stacks. */ +static _Thread_local _Alignas(32) float tl_qt[8192], tl_ot[8192], tl_st[G2M_ATTN_KB * G2M_ATTN_QB_MAX]; + +AVX2 static void g2x_media_attn(const g2m_attn_args *a) { + enum { KB = G2M_ATTN_KB }; + const int d = a->d, QB = g2m_attn_qblock(d), nq = a->nq; + float *qt = tl_qt, *ot = tl_ot, *st = tl_st; + _Alignas(32) float m[G2M_ATTN_QB_MAX], l[G2M_ATTN_QB_MAX], alpha[G2M_ATTN_QB_MAX], + lof[G2M_ATTN_QB_MAX], hif[G2M_ATTN_QB_MAX]; + const int np = (nq + 15) & ~15; /* query lanes, padded to whole 16-wide tiles */ + for (int dd = 0; dd < d; dd++) + for (int i = 0; i < np; i++) + qt[dd * QB + i] = i < nq ? a->q[(size_t)(a->q0 + i) * a->qstride + a->qcol + dd] * a->scale : 0.0f; + memset(ot, 0, (size_t)d * QB * sizeof(float)); + for (int i = 0; i < np; i++) { + m[i] = -INFINITY, l[i] = 0; + /* Padded lanes see every key (their q is zero, the results are dropped). */ + lof[i] = a->lo && i < nq ? (float)a->lo[i] : (float)a->klo; + hif[i] = a->hi && i < nq ? (float)a->hi[i] : (float)a->khi; + } + const bool masked = a->lo != NULL; + const float *kbase = a->k + a->kvcol, *vbase = a->v + a->kvcol; + + for (int b = a->klo; b < a->khi; b += KB) { + const int kn = a->khi - b < KB ? a->khi - b : KB; + /* S^T = K_blk . Q^T: 4 keys x 16 queries per register tile. */ + for (int j0 = 0; j0 < kn; j0 += 4) { + const float *kr[4]; + UNROLL for (int c = 0; c < 4; c++) + kr[c] = kbase + (size_t)(b + (j0 + c < kn ? j0 + c : kn - 1)) * a->kvstride; + for (int i0 = 0; i0 < np; i0 += 16) { + __m256 acc[4][2]; + UNROLL for (int c = 0; c < 4; c++) + acc[c][0] = acc[c][1] = _mm256_setzero_ps(); + const float *qp = qt + i0; + _Pragma("GCC unroll 4") for (int dd = 0; dd < d; dd++, qp += QB) { + __m256 a0 = _mm256_load_ps(qp), a1 = _mm256_load_ps(qp + 8); + UNROLL for (int c = 0; c < 4; c++) { + __m256 kb = _mm256_broadcast_ss(kr[c] + dd); + acc[c][0] = _mm256_fmadd_ps(kb, a0, acc[c][0]); + acc[c][1] = _mm256_fmadd_ps(kb, a1, acc[c][1]); + } + } + for (int c = 0; c < 4 && j0 + c < kn; c++) { + _mm256_store_ps(st + (j0 + c) * QB + i0, acc[c][0]); + _mm256_store_ps(st + (j0 + c) * QB + i0 + 8, acc[c][1]); + } + } + } + /* Mask, then online softmax, 8 query lanes at a time. */ + for (int i0 = 0; i0 < np; i0 += 8) { + __m256 mb = _mm256_set1_ps(-INFINITY); + if (masked) { + const __m256 lov = _mm256_load_ps(lof + i0), hiv = _mm256_load_ps(hif + i0); + for (int j = 0; j < kn; j++) { + __m256 key = _mm256_set1_ps((float)(b + j)), sv = _mm256_load_ps(st + j * QB + i0); + __m256 ok = _mm256_and_ps(_mm256_cmp_ps(key, lov, _CMP_GE_OQ), _mm256_cmp_ps(key, hiv, _CMP_LT_OQ)); + sv = _mm256_blendv_ps(_mm256_set1_ps(-INFINITY), sv, ok); + _mm256_store_ps(st + j * QB + i0, sv); + mb = _mm256_max_ps(mb, sv); + } + } else + for (int j = 0; j < kn; j++) + mb = _mm256_max_ps(mb, _mm256_load_ps(st + j * QB + i0)); + __m256 mo = _mm256_load_ps(m + i0), mn = _mm256_max_ps(mo, mb); + /* A lane with no visible key yet keeps m = -inf; subtract a finite + * floor so exp sees -inf - floor = -inf -> 0, never inf - inf. */ + __m256 ms = _mm256_max_ps(mn, _mm256_set1_ps(-3.0e38f)); + __m256 al = g2x_exp8(_mm256_sub_ps(mo, ms)), sum = _mm256_setzero_ps(); + for (int j = 0; j < kn; j++) { + __m256 pv = g2x_exp8(_mm256_sub_ps(_mm256_load_ps(st + j * QB + i0), ms)); + _mm256_store_ps(st + j * QB + i0, pv); + sum = _mm256_add_ps(sum, pv); + } + _mm256_store_ps(l + i0, _mm256_fmadd_ps(_mm256_load_ps(l + i0), al, sum)); + _mm256_store_ps(m + i0, mn); + _mm256_store_ps(alpha + i0, al); + } + /* O^T = O^T * alpha + V_blk^T . P^T: 4 dims x 16 queries per tile. */ + for (int d0 = 0; d0 < d; d0 += 4) + for (int i0 = 0; i0 < np; i0 += 16) { + __m256 al0 = _mm256_load_ps(alpha + i0), al1 = _mm256_load_ps(alpha + i0 + 8); + __m256 acc[4][2]; + UNROLL for (int c = 0; c < 4; c++) { + acc[c][0] = _mm256_mul_ps(_mm256_load_ps(ot + (d0 + c) * QB + i0), al0); + acc[c][1] = _mm256_mul_ps(_mm256_load_ps(ot + (d0 + c) * QB + i0 + 8), al1); + } + const float *vr = vbase + (size_t)b * a->kvstride + d0; + const float *pp = st + i0; + for (int j = 0; j < kn; j++, vr += a->kvstride, pp += QB) { + __m256 p0 = _mm256_load_ps(pp), p1 = _mm256_load_ps(pp + 8); + UNROLL for (int c = 0; c < 4; c++) { + __m256 vb = _mm256_broadcast_ss(vr + c); + acc[c][0] = _mm256_fmadd_ps(vb, p0, acc[c][0]); + acc[c][1] = _mm256_fmadd_ps(vb, p1, acc[c][1]); + } + } + UNROLL for (int c = 0; c < 4; c++) { + _mm256_store_ps(ot + (d0 + c) * QB + i0, acc[c][0]); + _mm256_store_ps(ot + (d0 + c) * QB + i0 + 8, acc[c][1]); + } + } + } + for (int i = 0; i < nq; i++) { + float inv = l[i] > 0 ? 1.0f / l[i] : 0.0f, *o = a->out + (size_t)(a->q0 + i) * a->ostride + a->ocol; + for (int dd = 0; dd < d; dd++) + o[dd] = ot[dd * QB + i] * inv; + } +} + +AVX2 static void g2x_media_quick_gelu_mul(float *g, const float *u, int n) { + const __m256 c = _mm256_set1_ps(-1.702f), one = _mm256_set1_ps(1.0f); + int z = 0; + for (; z + 8 <= n; z += 8) { + __m256 x = _mm256_loadu_ps(g + z); + __m256 y = _mm256_div_ps(x, _mm256_add_ps(one, g2x_exp8(_mm256_mul_ps(c, x)))); + _mm256_storeu_ps(g + z, u ? _mm256_mul_ps(y, _mm256_loadu_ps(u + z)) : y); + } + for (; z < n; z++) + g[z] = g[z] / (1.0f + expf(-1.702f * g[z])) * (u ? u[z] : 1.0f); +} + +AVX2 static void g2x_media_silu(float *y, const float *x, int n) { + const __m256 one = _mm256_set1_ps(1.0f), neg = _mm256_set1_ps(-1.0f); + int z = 0; + for (; z + 8 <= n; z += 8) { + __m256 a = _mm256_loadu_ps(x + z); + _mm256_storeu_ps(y + z, _mm256_div_ps(a, _mm256_add_ps(one, g2x_exp8(_mm256_mul_ps(neg, a))))); + } + for (; z < n; z++) + y[z] = x[z] / (1.0f + expf(-x[z])); +} + +AVX2 static void g2x_media_exp(float *x, int n) { + int z = 0; + for (; z + 8 <= n; z += 8) + _mm256_storeu_ps(x + z, g2x_exp8(_mm256_loadu_ps(x + z))); + for (; z < n; z++) + x[z] = expf(x[z]); +} +#endif + +void g2m_attn(const g2m_attn_args *a) { +#if defined(__x86_64__) || defined(_M_X64) + if (g2x_level() >= G2X_AVX2 && a->d % 16 == 0 && a->d <= 512) { + g2x_media_attn(a); + return; + } +#endif + g2m_attn_reference(a); +} + +void g2m_attn64(const float *q, const float *k, const float *v, float *out, size_t stride, size_t col, + int q0, int nq, int k0, int k1) { + g2m_attn_args a = {.q = q, .k = k, .v = v, .out = out, .qstride = stride, .kvstride = stride, + .ostride = stride, .qcol = col, .kvcol = col, .ocol = col, .d = 64, .q0 = q0, .nq = nq, + .klo = k0, .khi = k1, .scale = 1.0f}; + g2m_attn(&a); +} + +void g2m_quick_gelu_mul(float *g, const float *u, int n) { +#if defined(__x86_64__) || defined(_M_X64) + if (g2x_level() >= G2X_AVX2) { + g2x_media_quick_gelu_mul(g, u, n); + return; + } +#endif + for (int z = 0; z < n; z++) + g[z] = g[z] / (1.0f + expf(-1.702f * g[z])) * (u ? u[z] : 1.0f); +} + +void g2m_silu(float *y, const float *x, int n) { +#if defined(__x86_64__) || defined(_M_X64) + if (g2x_level() >= G2X_AVX2) { + g2x_media_silu(y, x, n); + return; + } +#endif + for (int z = 0; z < n; z++) + y[z] = x[z] / (1.0f + expf(-x[z])); +} + +void g2m_exp_inplace(float *x, int n) { +#if defined(__x86_64__) || defined(_M_X64) + if (g2x_level() >= G2X_AVX2) { + g2x_media_exp(x, n); + return; + } +#endif + for (int z = 0; z < n; z++) + x[z] = expf(x[z]); +} diff --git a/src/gemma2_media_x86.h b/src/gemma2_media_x86.h new file mode 100644 index 0000000..4acce7b --- /dev/null +++ b/src/gemma2_media_x86.h @@ -0,0 +1,50 @@ +/* SIMD kernels shared by the EmbeddingGemma 2 CPU text graph and media towers + * (vision/audio): tiled online-softmax attention, QuickGELU, SiLU, exp. x86-64 + * AVX2+FMA bodies are selected at run time (g2x_level) with portable C + * fallbacks, so generic builds stay generic. */ +#ifndef EI_GEMMA2_MEDIA_X86_H +#define EI_GEMMA2_MEDIA_X86_H +#include +#include + +/* Key block of the tiled kernels. Query blocks depend on the head width and + * kernel variant: callers must size tasks with g2m_attn_qblock(d), which never + * exceeds G2M_ATTN_QB_MAX. */ +#ifndef G2M_KB +#define G2M_KB 64 /* compile-time key block (tuning experiments: -DG2M_KB=32|128) */ +#endif +enum { G2M_ATTN_KB = G2M_KB, G2M_ATTN_QB_MAX = 128 }; + +/* Online-softmax attention for one head and a block of queries [q0, q0 + nq), + * nq <= g2m_attn_qblock(d). Head width d: multiple of 16, <= 512 (fp32, any d + * via the scalar path). Rows of q / k,v / out are qstride / kvstride / ostride + * floats apart and the head starts at column qcol / kvcol / ocol. Query i sees + * keys [lo[i], hi[i]) when lo/hi are given (sliding windows, packed + * sequences), else [klo, khi); klo/khi must bound every lo/hi. Scores are + * scaled by `scale`. Masked keys contribute exactly zero. */ +#include +typedef struct { + const float *q, *k, *v; + float *out; + size_t qstride, kvstride, ostride, qcol, kvcol, ocol; + int d, q0, nq; + const int32_t *lo, *hi; + int klo, khi; + float scale; +} g2m_attn_args; +int g2m_attn_qblock(int d); +void g2m_attn(const g2m_attn_args *a); +/* Scalar implementation (the non-AVX2 fallback; used by tests as the reference). */ +void g2m_attn_reference(const g2m_attn_args *a); +/* Full bidirectional softmax attention, scale 1, for one 64-wide head: + * rows of q/k/v/out are `stride` floats apart; the head starts at column + * `col`. Keys [k0, k1) are visible to every query of the task. */ +void g2m_attn64(const float *q, const float *k, const float *v, float *out, size_t stride, size_t col, + int q0, int nq, int k0, int k1); +/* g[i] = act(g[i]) * u[i] (u NULL: act only); act = x * sigmoid(1.702 x) (QuickGELU). */ +void g2m_quick_gelu_mul(float *g, const float *u, int n); +/* y[i] = x[i] / (1 + exp(-x[i])) (SiLU), in place when y == x. */ +void g2m_silu(float *y, const float *x, int n); +/* The attention kernel's vector exp, in place (accuracy checked by test_gemma2_media). */ +void g2m_exp_inplace(float *x, int n); +#endif diff --git a/src/gemma2_release.h b/src/gemma2_release.h new file mode 100644 index 0000000..dbf2a91 --- /dev/null +++ b/src/gemma2_release.h @@ -0,0 +1,38 @@ +/* Drop the resident pages of a weight tensor once a kernel holds + * a repacked copy. The GGUF is mapped PROT_READ | MAP_PRIVATE, so its pages are + * never dirty: MADV_DONTNEED unmaps them (they are re-read from the file if + * anything touches them again) and the process keeps one copy of each repacked + * weight instead of two. Only whole pages inside the tensor are released. + * EI_GEMMA2_CPU_RELEASE_WEIGHTS=0 keeps the source pages mapped. */ +#ifndef EI_GEMMA2_RELEASE_H +#define EI_GEMMA2_RELEASE_H +#include +#include +#include +#include +#include "gguf.h" +#if defined(__linux__) || defined(__APPLE__) +#include +#include +#endif + +static inline void g2_release_source(const void *data, size_t bytes) { +#if defined(__linux__) || defined(__APPLE__) + static int on = -1; + if (on < 0) { + const char *e = getenv("EI_GEMMA2_CPU_RELEASE_WEIGHTS"); + on = !(e && !strcmp(e, "0")); + } + /* Only pages of a read-only GGUF mapping: on anonymous memory (weights built + * in memory, e.g. by tests) MADV_DONTNEED would zero the data. */ + if (!on || !ei_gguf_in_mapping(data, bytes)) + return; + const uintptr_t page = (uintptr_t)sysconf(_SC_PAGESIZE); + uintptr_t lo = ((uintptr_t)data + page - 1) & ~(page - 1), hi = ((uintptr_t)data + bytes) & ~(page - 1); + if (hi > lo) + madvise((void *)lo, hi - lo, MADV_DONTNEED); +#else + (void)data, (void)bytes; +#endif +} +#endif diff --git a/src/gemma2_sgemm_x86.c b/src/gemma2_sgemm_x86.c new file mode 100644 index 0000000..7216ad3 --- /dev/null +++ b/src/gemma2_sgemm_x86.c @@ -0,0 +1,141 @@ +#if defined(__linux__) +#define _DEFAULT_SOURCE /* madvise under -std=c11 (gemma2_release.h) */ +#endif +/* Packed fp32 GEMM, y[n][m] = x[n][k] . w[m][k]^T. + * + * The row float-weight GEMM (gemma2_fast.c gemm_f, f_rows) converts four weight + * rows to fp32 per task on every call and computes 4x4 dot tiles that end in + * 16 horizontal reductions. Here the weight is expanded once (BF16/F16 -> fp32 + * is exact) into panels of 16 output columns, wp[m/16][k][16], and a 6-token x + * 16-output register tile (12 ymm accumulators) streams one panel row per k: + * 2 loads + 6 broadcasts per 12 FMAs, no reductions. fp32 throughout; only + * the summation order differs from the reference. */ +#include "gemma2_sgemm_x86.h" +#include "gemma2_fast_x86.h" +#include "gemma2_release.h" +#include "quants.h" + +#include +#include + +typedef struct { + const void *key; + int m, k; + float *p; /* [m/16][k][16] */ +} g2s_weight; + +int g2s_m(const void *w) { return ((const g2s_weight *)w)->m; } +int g2s_k(const void *w) { return ((const g2s_weight *)w)->k; } + +#if defined(__x86_64__) || defined(_M_X64) +enum { G2S_CACHE = 64 }; +static g2s_weight g2s_weights[G2S_CACHE]; +static int g2s_count; +#include +#define AVX2 __attribute__((target("avx2,fma,f16c"))) +#define UNROLL _Pragma("GCC unroll 8") + +int g2s_enabled(void) { + static int cached = -1; + if (cached < 0) { + const char *flag = getenv("EI_GEMMA2_CPU_SGEMM"); /* default on where supported */ + cached = (!flag || strcmp(flag, "0")) && g2x_level() >= G2X_AVX2; + } + return cached; +} + +static float element(const void *data, int type, size_t i) { + if (type == EI_T_F32) + return ((const float *)data)[i]; + if (type == EI_T_F16) + return ei_fp16_to_fp32(((const ei_fp16 *)data)[i]); + uint32_t bits = (uint32_t)((const uint16_t *)data)[i] << 16; /* BF16 */ + float f; + memcpy(&f, &bits, 4); + return f; +} + +typedef struct { + float *p; + const void *data; + int type, k; +} g2s_pack_job; +static void g2s_pack_panels(void *opaque, int32_t begin, int32_t end) { + const g2s_pack_job *j = opaque; + for (int pn = begin; pn < end; pn++) + for (int z = 0; z < j->k; z++) + for (int c = 0; c < 16; c++) + j->p[((size_t)pn * j->k + z) * 16 + c] = element(j->data, j->type, (size_t)(pn * 16 + c) * j->k + z); +} + +const void *g2s_weight_of(const void *key, const void *data, int type, int m, int k, ei_thread_pool *pool) { + for (int i = 0; i < g2s_count; i++) + if (g2s_weights[i].key == key) + return g2s_weights[i].m == m && g2s_weights[i].k == k ? &g2s_weights[i] : NULL; + if (g2s_count == G2S_CACHE || m % 16 || (type != EI_T_F32 && type != EI_T_F16 && type != EI_T_BF16)) + return NULL; + g2s_weight *w = &g2s_weights[g2s_count]; + w->p = ei_aligned_alloc(64, (size_t)m * k * sizeof(float)); + g2s_pack_job job = {w->p, data, type, k}; + ei_parallel_for(pool, m / 16, 4, g2s_pack_panels, &job); /* on the caller's pool (first use) */ + g2_release_source(data, (size_t)m * k * (type == EI_T_F32 ? 4 : 2)); + w->key = key; + w->m = m; + w->k = k; + g2s_count++; + return w; +} + +/* g2x_ prefix: scripts/check-avx2-isolation.py's dispatched set. */ +#define G2S_TILE(NT) \ + AVX2 static void g2x_sgemm_tile##NT(const float *panel, const float *x, int k, float *y, int m) { \ + __m256 acc[NT][2]; \ + UNROLL for (int t = 0; t < NT; t++) acc[t][0] = acc[t][1] = _mm256_setzero_ps(); \ + _Pragma("GCC unroll 2") for (int z = 0; z < k; z++) { \ + __m256 b0 = _mm256_load_ps(panel + (size_t)z * 16), b1 = _mm256_load_ps(panel + (size_t)z * 16 + 8);\ + UNROLL for (int t = 0; t < NT; t++) { \ + __m256 a = _mm256_broadcast_ss(x + (size_t)t * k + z); \ + acc[t][0] = _mm256_fmadd_ps(a, b0, acc[t][0]); \ + acc[t][1] = _mm256_fmadd_ps(a, b1, acc[t][1]); \ + } \ + } \ + UNROLL for (int t = 0; t < NT; t++) { \ + _mm256_storeu_ps(y + (size_t)t * m, acc[t][0]); \ + _mm256_storeu_ps(y + (size_t)t * m + 8, acc[t][1]); \ + } \ + } +G2S_TILE(1) +G2S_TILE(2) +G2S_TILE(3) +G2S_TILE(4) +G2S_TILE(5) +G2S_TILE(6) + +void g2s_block(const void *wp, const float *x, int t0, int t1, int p0, int np, float *y) { + const g2s_weight *w = wp; + for (int pn = p0; pn < p0 + np; pn++) { + const float *panel = w->p + (size_t)pn * w->k * 16; + for (int t = t0; t < t1; t += 6) { + const float *xt = x + (size_t)t * w->k; + float *yt = y + (size_t)t * w->m + pn * 16; + switch (t1 - t < 6 ? t1 - t : 6) { + case 1: g2x_sgemm_tile1(panel, xt, w->k, yt, w->m); break; + case 2: g2x_sgemm_tile2(panel, xt, w->k, yt, w->m); break; + case 3: g2x_sgemm_tile3(panel, xt, w->k, yt, w->m); break; + case 4: g2x_sgemm_tile4(panel, xt, w->k, yt, w->m); break; + case 5: g2x_sgemm_tile5(panel, xt, w->k, yt, w->m); break; + default: g2x_sgemm_tile6(panel, xt, w->k, yt, w->m); break; + } + } + } +} +#else +int g2s_enabled(void) { return 0; } +const void *g2s_weight_of(const void *key, const void *data, int type, int m, int k, ei_thread_pool *pool) { + (void)key, (void)data, (void)type, (void)m, (void)k, (void)pool; + return NULL; +} +void g2s_block(const void *w, const float *x, int t0, int t1, int p0, int np, float *y) { + (void)w, (void)x, (void)t0, (void)t1, (void)p0, (void)np, (void)y; +} +#endif diff --git a/src/gemma2_sgemm_x86.h b/src/gemma2_sgemm_x86.h new file mode 100644 index 0000000..b44cc23 --- /dev/null +++ b/src/gemma2_sgemm_x86.h @@ -0,0 +1,19 @@ +/* Packed fp32 GEMM for the float-weight projections of the + * EmbeddingGemma 2 CPU graphs (BF16 per-layer-input projection, F32 vision + * patch embedding and audio input projection). x86 AVX2+FMA only. */ +#ifndef EI_GEMMA2_SGEMM_X86_H +#define EI_GEMMA2_SGEMM_X86_H +#include "parallel.h" +#include +#include + +/* 1 when the kernel can run here and EI_GEMMA2_CPU_SGEMM != 0. */ +int g2s_enabled(void); +/* Weights [m][k] of the given ei_tensor_type (F32, F16 or BF16), expanded once + * (exactly) into 16-column fp32 panels and cached by `key`. NULL if m % 16. */ +const void *g2s_weight_of(const void *key, const void *data, int type, int m, int k, ei_thread_pool *pool); +/* y[t][m0 .. m0 + 16*np) for tokens [t0, t1), panels [p0, p0 + np). */ +void g2s_block(const void *w, const float *x, int t0, int t1, int p0, int np, float *y); +int g2s_m(const void *w); +int g2s_k(const void *w); +#endif diff --git a/src/gemma2_vision_fast.c b/src/gemma2_vision_fast.c new file mode 100644 index 0000000..bcbc786 --- /dev/null +++ b/src/gemma2_vision_fast.c @@ -0,0 +1,370 @@ +/* SIMD + threaded CPU vision tower for EmbeddingGemma 2. + * + * The reference CPU vision tower is the scalar graph + * (gemma2.c: per-element GEMM loops, scalar attention with a calloc per job, + * single-threaded norms/RoPE/activations/quantization). This path keeps the + * reference math -- Q8_0 activations for Q8_0 weights, fp32 attention and + * norms, the fp16-rounded patch input, the selected GELU variant -- and + * changes only how it is executed: + * - Q8_0 GEMMs on the fast text graph's x86 VNNI/AVX2 (or NEON) tiles, + * - g2m_attn64: tiled online-softmax attention for the 64-wide heads, + * - per-token fused row passes (norm+quantize, q/k/v norms + axial RoPE from + * tables, GELU+quantize, post-norm+residual), threaded over tokens, + * - the 3072-wide FFN intermediate processed in token chunks, so its buffers + * are O(chunk) instead of O(patches), + * - a persistent workspace instead of nine callocs per image. + * EI_GEMMA2_CPU_MEDIA_FAST=0 selects the reference graph. */ +#include "gemma2.h" +#include "gemma2_fast_x86.h" +#include "gemma2_media_x86.h" +#include "kernels.h" +#include +#include +#include +#include + +enum { VD = 768, VF = 3072, VHEADS = 12, FFN_CHUNK = 256 }; + +/* EI_GEMMA2_CPU_PROFILE=1: per-stage wall time of each vision call on stderr. */ +#include +enum { VP_PATCH, VP_GEMM_QKV, VP_ROWS, VP_ATTN, VP_GEMM_O, VP_GEMM_FFN, VP_POOL, VP_N }; +static const char *vp_names[VP_N] = {"patch", "gemm_qkv", "rows", "attention", "gemm_out", "gemm_ffn", "pool_proj"}; +static double vp_ms[VP_N]; +static int vp_on(void) { + static int on = -1; + if (on < 0) { + const char *e = getenv("EI_GEMMA2_CPU_PROFILE"); + on = e && strcmp(e, "0") != 0; + } + return on; +} +static double vp_now(void) { + struct timespec ts; + clock_gettime(CLOCK_MONOTONIC, &ts); + return ts.tv_sec * 1e3 + ts.tv_nsec * 1e-6; +} +#define VP(i, ...) \ + do { \ + double vp_t0_ = vp_on() ? vp_now() : 0; \ + __VA_ARGS__; \ + if (vp_on()) \ + vp_ms[i] += vp_now() - vp_t0_; \ + } while (0) + +bool ei_gemma2_cpu_media_fast(void) { + static int cached = -1; + if (cached < 0) { + const char *flag = getenv("EI_GEMMA2_CPU_MEDIA_FAST"); /* default on; "0" = reference */ + cached = !flag || strcmp(flag, "0"); + } + return cached; +} + +static const ei_tensor *tensor_of(const ei_gguf *g, const char *name) { + return ei_gguf_tensor(g, name, true); +} +static const float *f32_of(const ei_gguf *g, const char *name) { + const ei_tensor *t = tensor_of(g, name); + if (t->type != EI_T_F32) + ei_die("EmbeddingGemma 2 vision: %s must be F32", name); + return t->data; +} + +/* ---- workspace ------------------------------------------------------------ */ + +typedef struct { + size_t cap; /* patches */ + float *x, *h, *q, *k, *v; /* [cap][768]; v doubles as the attention output */ + float *up, *gate; /* [FFN_CHUNK][3072] */ + ei_block_q8_0 *xq; /* [max(cap*24, FFN_CHUNK*96)] */ + float *rope; /* [positions][32]: 16 cos then 16 sin */ + int rope_positions; +} vision_ws; + +/* One workspace per process, grown to the largest image seen and kept: at + * the 1120-soft-token cap (10080 patches) the five [patches][768] planes are + * about 150 MB, which a server reaching that size holds anyway. */ +static pthread_mutex_t vision_lock = PTHREAD_MUTEX_INITIALIZER; +static vision_ws vws; +static ei_gemma2_fast *gemm_ctx; + +static void reserve(vision_ws *w, size_t n) { + if (w->cap >= n) + return; + free(w->x), free(w->h), free(w->q), free(w->k), free(w->v), free(w->xq); + size_t row = (size_t)VD * sizeof(float); + w->x = ei_xmalloc(n * row); + w->h = ei_xmalloc(n * row); + w->q = ei_xmalloc(n * row); + w->k = ei_xmalloc(n * row); + w->v = ei_xmalloc(n * row); + size_t blocks = n * VD / 32 > (size_t)FFN_CHUNK * VF / 32 ? n * VD / 32 : (size_t)FFN_CHUNK * VF / 32; + w->xq = ei_xmalloc(blocks * sizeof(ei_block_q8_0)); + if (!w->up) { + w->up = ei_xmalloc((size_t)FFN_CHUNK * VF * sizeof(float)); + w->gate = ei_xmalloc((size_t)FFN_CHUNK * VF * sizeof(float)); + } + w->cap = n; +} + +/* theta = p * 100^(-2z/32), the reference's exact expression, z = 0..15. */ +static void rope_table(vision_ws *w, int positions) { + if (w->rope_positions >= positions) + return; + free(w->rope); + w->rope = ei_xmalloc((size_t)positions * 32 * sizeof(float)); + for (int p = 0; p < positions; p++) + for (int z = 0; z < 16; z++) { + float theta = p * powf(100.0f, -2.0f * z / 32); + w->rope[(size_t)p * 32 + z] = cosf(theta); + w->rope[(size_t)p * 32 + 16 + z] = sinf(theta); + } + w->rope_positions = positions; +} + +/* ---- row passes ------------------------------------------------------------- */ + +typedef enum { ROW_NORM_QUANT, ROW_QKV_POST, ROW_V_HEADS, ROW_QUANT, ROW_POST_ADD, ROW_GELU_QUANT, ROW_POOL } row_op; +typedef struct { + row_op op; + vision_ws *w; + const float *nw, *qn, *kn; /* norm weights */ + float *src, *dst; + int cols, out_cols, t0; /* ROW_GELU_QUANT: first token of the chunk */ + int n; /* ROW_QKV_POST / ROW_V_HEADS: tokens (head-major planes) */ + bool tanh_gelu; +} row_job; + +static void rms(float *y, const float *x, int d, const float *w) { + float scale = 1.0f / sqrtf(ei_gemma2_fast_sum_squares(x, d) / d + 1e-6f); + if (w) + for (int z = 0; z < d; z++) + y[z] = x[z] * (scale * w[z]); + else + for (int z = 0; z < d; z++) + y[z] = x[z] * scale; +} +static void rope_head(float *r, const float *cx, const float *cy) { + for (int axis = 0; axis < 2; axis++) { + const float *c = axis ? cy : cx, *s = c + 16; + float *h = r + axis * 32; + for (int z = 0; z < 16; z++) { + float a = h[z], b = h[z + 16]; + h[z] = a * c[z] - b * s[z]; + h[z + 16] = a * s[z] + b * c[z]; + } + } +} +static void row_tokens(void *opaque, int32_t begin, int32_t end) { + row_job *j = opaque; + vision_ws *w = j->w; + float tmp[VD]; + for (int t = begin; t < end; t++) { + switch (j->op) { + case ROW_NORM_QUANT: /* xq = q8(rms(x) * w) */ + rms(tmp, w->x + (size_t)t * VD, VD, j->nw); + ei_gemma2_fast_quant_row(tmp, w->xq + (size_t)t * (VD / 32), VD); + break; + case ROW_QKV_POST: { /* q in place; k -> h as [head][n][64] (see attn_tasks) */ + const float *cx = w->rope + (size_t)(t % j->cols) * 32, *cy = w->rope + (size_t)(t / j->cols) * 32; + for (int hd = 0; hd < VHEADS; hd++) { + float *qr = w->q + (size_t)t * VD + hd * 64, *kh = w->h + ((size_t)hd * j->n + t) * 64; + rms(qr, qr, 64, j->qn); + rms(kh, w->k + (size_t)t * VD + hd * 64, 64, j->kn); + rope_head(qr, cx, cy); + rope_head(kh, cx, cy); + } + break; + } + case ROW_V_HEADS: /* v -> the k buffer (free after ROW_QKV_POST) as [head][n][64] */ + for (int hd = 0; hd < VHEADS; hd++) + rms(w->k + ((size_t)hd * j->n + t) * 64, w->v + (size_t)t * VD + hd * 64, 64, NULL); + break; + case ROW_QUANT: + ei_gemma2_fast_quant_row(j->src + (size_t)t * VD, w->xq + (size_t)t * (VD / 32), VD); + break; + case ROW_POST_ADD: { /* x += rms(h) * w */ + float *x = w->x + (size_t)t * VD; + rms(tmp, w->h + (size_t)t * VD, VD, j->nw); + for (int z = 0; z < VD; z++) + x[z] += tmp[z]; + break; + } + case ROW_GELU_QUANT: { /* chunk-local row t: gate = act(gate) * up; q8 */ + float *g = w->gate + (size_t)t * VF; + const float *u = w->up + (size_t)t * VF; + if (j->tanh_gelu) + ei_gemma2_fast_gelu_mul(g, u, VF); + else + g2m_quick_gelu_mul(g, u, VF); + ei_gemma2_fast_quant_row(g, w->xq + (size_t)t * (VF / 32), VF); + break; + } + case ROW_POOL: { /* 3x3 mean * sqrt(768), unweighted rms, q8 (reference order) */ + float *o = tmp; + memset(o, 0, sizeof tmp); + for (int yy = 0; yy < 3; yy++) + for (int xx = 0; xx < 3; xx++) { + const float *s = w->x + ((size_t)(t / j->out_cols * 3 + yy) * j->cols + t % j->out_cols * 3 + xx) * VD; + for (int z = 0; z < VD; z++) + o[z] += s[z] / 9.0f * sqrtf(768.0f); + } + rms(o, o, VD, NULL); + ei_gemma2_fast_quant_row(o, w->xq + (size_t)t * (VD / 32), VD); + break; + } + } + } +} +static void rows(ei_thread_pool *pool, int n, row_job job) { ei_parallel_for(pool, n, 4, row_tokens, &job); } + +/* ---- patch embedding ------------------------------------------------------- */ + +typedef struct { + const uint8_t *rgb; + float *patches; + int width, cols; +} patch_job; +static void patch_tokens(void *opaque, int32_t begin, int32_t end) { + patch_job *j = opaque; + for (int t = begin; t < end; t++) + for (int c = 0; c < 3; c++) + for (int y = 0; y < 16; y++) + for (int z = 0; z < 16; z++) { + float p = 2.0f * j->rgb[((size_t)(t / j->cols * 16 + y) * j->width + t % j->cols * 16 + z) * 3 + c] / 255.0f - 1; + /* The reference convolution's fp16 input rounding. */ + j->patches[(size_t)t * VD + c * 256 + y * 16 + z] = ei_fp16_to_fp32(ei_fp32_to_fp16(p)); + } +} +typedef struct { + float *x; + const float *pos; + int cols; +} pos_job; +static void pos_tokens(void *opaque, int32_t begin, int32_t end) { + pos_job *j = opaque; + for (int t = begin; t < end; t++) { + const float *px = j->pos + (size_t)(t % j->cols) * VD, *py = j->pos + ((size_t)10240 + t / j->cols) * VD; + float *x = j->x + (size_t)t * VD; + for (int z = 0; z < VD; z++) + x[z] += px[z] + py[z]; + } +} + +/* ---- attention ----------------------------------------------------------------- */ + +typedef struct { + const vision_ws *w; + int n, blocks; +} attn_job; +/* Q and the output (into v, free once V moved) are token-major [n][768]; K (in h) and V (in k) are + * head-major [12][n][64], so the K/V stream of a query block is contiguous: + * read in place from the token-major buffers, every key of a head sat on its + * own 4 KB page (3 KB row stride) and the per-head stream at 10080 patches + * (5.2 MB) defeated the L2 prefetchers. */ +static void attn_tasks(void *opaque, int32_t begin, int32_t end) { + attn_job *j = opaque; + for (int task = begin; task < end; task++) { + const int qb = g2m_attn_qblock(64); + int hd = task / j->blocks, q0 = task % j->blocks * qb; + int nq = j->n - q0 < qb ? j->n - q0 : qb; + const size_t plane = (size_t)hd * j->n * 64; + g2m_attn_args a = {.q = j->w->q, .k = j->w->h + plane, .v = j->w->k + plane, .out = j->w->v, + .qstride = VD, .kvstride = 64, .ostride = VD, .qcol = (size_t)hd * 64, .kvcol = 0, + .ocol = (size_t)hd * 64, .d = 64, .q0 = q0, .nq = nq, .klo = 0, .khi = j->n, + .scale = 1.0f}; + g2m_attn(&a); + } +} + +/* ---- tower --------------------------------------------------------------------- */ + +static void create_ctx(void) { gemm_ctx = ei_gemma2_fast_create(); } + +bool ei_gemma2_vision_fast(const ei_gguf *g, ei_thread_pool *pool, const uint8_t *rgb, int width, + int height, float **output, size_t *tokens, char *err, size_t err_len) { + if (width < 48 || height < 48 || width % 48 || height % 48 || + (size_t)width * height > 1120 * 48 * 48) { + snprintf(err, err_len, "vision input must be aligned to 48 pixels and at most 1120 soft tokens"); + return false; + } + static pthread_once_t once = PTHREAD_ONCE_INIT; + pthread_once(&once, create_ctx); + const int cols = width / 16, prow = height / 16, n = cols * prow; + const bool tanh_gelu = ei_gemma2_vision_tanh_gelu(g); + + pthread_mutex_lock(&vision_lock); + ei_cpu_binding binding; /* hybrid pools: the caller joins on a P-core */ + ei_thread_pool_bind_caller(pool, &binding); + vision_ws *w = &vws; + reserve(w, (size_t)n); + rope_table(w, cols > prow ? cols : prow); + + /* Patch embedding (F32 weights) + learned x/y positions. Patches stage in q. */ + patch_job pj = {rgb, w->q, width, cols}; + ei_tensor patch = *tensor_of(g, "v.patch_embd.weight"); + patch.ne[0] = VD; + patch.ne[1] = VD; + pos_job po = {w->x, f32_of(g, "v.position_embd.weight"), cols}; + VP(VP_PATCH, ei_parallel_for(pool, n, 16, patch_tokens, &pj); + ei_gemma2_fast_gemm_f(pool, &patch, w->q, n, w->x); + ei_parallel_for(pool, n, 16, pos_tokens, &po)); + + attn_job aj = {w, n, (n + g2m_attn_qblock(64) - 1) / g2m_attn_qblock(64)}; + for (int l = 0; l < 16; l++) { + char name[128]; +#define V(suffix) (snprintf(name, sizeof name, "v.blk.%d." suffix ".weight", l), tensor_of(g, name)) +#define VF32(suffix) (snprintf(name, sizeof name, "v.blk.%d." suffix ".weight", l), f32_of(g, name)) + if (l) + ei_gemma2_cpu_yield(); /* between layers: waiting text waves may run */ + VP(VP_ROWS, rows(pool, n, (row_job){.op = ROW_NORM_QUANT, .w = w, .nw = VF32("ln1")})); + VP(VP_GEMM_QKV, ei_gemma2_fast_gemm_q8(gemm_ctx, pool, V("attn_q"), w->xq, n, w->q); + ei_gemma2_fast_gemm_q8(gemm_ctx, pool, V("attn_k"), w->xq, n, w->k); + ei_gemma2_fast_gemm_q8(gemm_ctx, pool, V("attn_v"), w->xq, n, w->v)); + const float *qn = VF32("attn_q_norm"), *kn = VF32("attn_k_norm"); + VP(VP_ROWS, rows(pool, n, (row_job){.op = ROW_QKV_POST, .w = w, .qn = qn, .kn = kn, .cols = cols, .n = n}); + rows(pool, n, (row_job){.op = ROW_V_HEADS, .w = w, .n = n})); + VP(VP_ATTN, ei_parallel_for(pool, VHEADS * aj.blocks, 1, attn_tasks, &aj)); + ei_gemma2_cpu_yield(); + VP(VP_ROWS, rows(pool, n, (row_job){.op = ROW_QUANT, .w = w, .src = w->v})); + VP(VP_GEMM_O, ei_gemma2_fast_gemm_q8(gemm_ctx, pool, V("attn_out"), w->xq, n, w->h)); + VP(VP_ROWS, rows(pool, n, (row_job){.op = ROW_POST_ADD, .w = w, .nw = VF32("attn_post_norm")})); + /* ln2 activations for every token first (q buffer is free now), then the + * FFN in chunks: up/gate/act/down per FFN_CHUNK tokens. */ + VP(VP_ROWS, rows(pool, n, (row_job){.op = ROW_NORM_QUANT, .w = w, .nw = VF32("ln2")})); + memcpy(w->q, w->xq, (size_t)n * (VD / 32) * sizeof(ei_block_q8_0)); + const ei_block_q8_0 *ln2q = (const ei_block_q8_0 *)w->q; + const ei_tensor *up = V("ffn_up"), *gate = V("ffn_gate"), *down = V("ffn_down"); + for (int c0 = 0; c0 < n; c0 += FFN_CHUNK) { + int cn = n - c0 < FFN_CHUNK ? n - c0 : FFN_CHUNK; + if (c0) + ei_gemma2_cpu_yield(); + const ei_block_q8_0 *xin = ln2q + (size_t)c0 * (VD / 32); + VP(VP_GEMM_FFN, ei_gemma2_fast_gemm_q8(gemm_ctx, pool, up, xin, cn, w->up); + ei_gemma2_fast_gemm_q8(gemm_ctx, pool, gate, xin, cn, w->gate)); + VP(VP_ROWS, rows(pool, cn, (row_job){.op = ROW_GELU_QUANT, .w = w, .tanh_gelu = tanh_gelu})); + VP(VP_GEMM_FFN, ei_gemma2_fast_gemm_q8(gemm_ctx, pool, down, w->xq, cn, w->h + (size_t)c0 * VD)); + } + VP(VP_ROWS, rows(pool, n, (row_job){.op = ROW_POST_ADD, .w = w, .nw = VF32("ffn_post_norm")})); +#undef V +#undef VF32 + } + const int out_cols = cols / 3, nt = out_cols * (prow / 3); + *output = ei_xmalloc((size_t)nt * 512 * sizeof(float)); + *tokens = (size_t)nt; + VP(VP_POOL, rows(pool, nt, (row_job){.op = ROW_POOL, .w = w, .cols = cols, .out_cols = out_cols}); + ei_gemma2_fast_gemm_q8(gemm_ctx, pool, tensor_of(g, "mm.input_projection.weight"), w->xq, nt, *output)); + if (vp_on()) { + double total = 0; + for (int i = 0; i < VP_N; i++) + total += vp_ms[i]; + fprintf(stderr, "vision-fast patches=%d total=%.1f ms:", n, total); + for (int i = 0; i < VP_N; i++) + fprintf(stderr, " %s=%.1f", vp_names[i], vp_ms[i]); + fprintf(stderr, "\n"); + memset(vp_ms, 0, sizeof vp_ms); + } + ei_thread_pool_unbind_caller(&binding); + pthread_mutex_unlock(&vision_lock); + return true; +} diff --git a/src/gemma2_vnni256.c b/src/gemma2_vnni256.c new file mode 100644 index 0000000..78e2e7f --- /dev/null +++ b/src/gemma2_vnni256.c @@ -0,0 +1,196 @@ +#if defined(__linux__) +#define _DEFAULT_SOURCE /* madvise under -std=c11 (gemma2_release.h) */ +#endif +/* Repacked 256-bit AVX-VNNI Q8_0 x Q8_0 GEMM. + * + * The generic x86 Q8 tile (gemma2_fast_x86.c, q8_tile_avxvnni) computes each + * (row, token, 32-block) dot with abs(w), sign(x, w), a zero-seeded vpdpbusd, + * a conversion, a scalar scale product and an FMA, then reduces 8 lanes per + * output at the end: about 6 vector ops per 32 MACs. Its repacked kernel + * (g2x_gemm512_tile) needs AVX-512. + * + * This kernel repacks each weight once into 8-row groups so that one 32-byte + * load holds 4 consecutive k-bytes of 8 different output rows: + * w8[g][block][j][r][4] = w[8g + r][32 block + 4j .. 4j + 3] (j = 0..7) + * Per (group, block) the 8 rows' fp16 Q8_0 scales and int16 weight sums are + * kept (36 bytes per 32 weights, vs 34 for Q8_0 itself). + * Activations are the u8 = q + 128 rows of g2x_acts. For each token the 4 + * k-bytes of step j are broadcast to all lanes and one vpdpbusd advances 8 + * output rows at once; seeding the int32 accumulator with -128 * sum(w) per + * (row, block) cancels the +128 bias exactly. A block's 8 rows are then + * scaled with one convert, one multiply and one FMA (per-row weight scales + * times the token's block scale), and the 8 lanes are the 8 contiguous + * outputs -- no horizontal reduction. Per 256 MACs: 8 vpdpbusd + 3 ops. + * Integer dot products are exact, so results equal the reference Q8_0 dot up + * to fp32 summation order. */ +#include "gemma2_vnni256.h" +#include "gemma2_fast_x86.h" +#include "gemma2_release.h" + +#include +#include + +typedef struct { + const void *key; + int m, k; + int8_t *q; /* [m/8][k/32][8 steps][8 rows][4] */ + ei_fp16 *s; /* [m/8][k/32][8 rows], the Q8_0 block scales (fp16, exact) */ + int16_t *c; /* [m/8][k/32][8 rows] = sum(w) in [-4064, 4064]; x -128 at use */ +} g2v_weight; + +#if defined(__x86_64__) || defined(_M_X64) +#include + +enum { G2V_CACHE = 1024 }; +static g2v_weight g2v_weights[G2V_CACHE]; +static int g2v_count; + +int g2v_enabled(void) { + static int cached = -1; + if (cached < 0) { + const char *flag = getenv("EI_GEMMA2_CPU_VNNI256"); /* default on where supported */ + __builtin_cpu_init(); + cached = (!flag || strcmp(flag, "0")) && g2x_level() == G2X_AVXVNNI; /* level implies AVX-VNNI */ + } + return cached; +} + +/* Repack 8-row groups [begin, end): runs on the caller's pool (first use of a tensor). */ +typedef struct { + g2v_weight *w; + const ei_block_q8_0 *blocks; + int kb; +} g2v_repack_job; +static void g2v_repack_groups(void *opaque, int32_t begin, int32_t end) { + const g2v_repack_job *j = opaque; + const int kb = j->kb; + for (int g = begin; g < end; g++) + for (int p = 0; p < kb; p++) { + int8_t *dst = j->w->q + ((size_t)g * kb + p) * 256; + for (int r = 0; r < 8; r++) { + const ei_block_q8_0 *b = j->blocks + (size_t)(8 * g + r) * kb + p; + int32_t sum = 0; + for (int i = 0; i < 8; i++) + memcpy(dst + i * 32 + r * 4, b->qs + 4 * i, 4); + for (int z = 0; z < 32; z++) + sum += b->qs[z]; + j->w->s[((size_t)g * kb + p) * 8 + r] = b->d; + j->w->c[((size_t)g * kb + p) * 8 + r] = (int16_t)sum; + } + } +} + +const void *g2v_weight_of(const void *key, const ei_block_q8_0 *blocks, int m, int k, ei_thread_pool *pool) { + for (int i = 0; i < g2v_count; i++) + if (g2v_weights[i].key == key) { + if (g2v_weights[i].m != m || g2v_weights[i].k != k) + return NULL; /* a different tensor reused the address: never serve a stale repack */ + return &g2v_weights[i]; + } + if (g2v_count == G2V_CACHE || m % 8 || k % 32) + return NULL; + const int kb = k / 32; + g2v_weight *w = &g2v_weights[g2v_count]; + w->q = ei_aligned_alloc(64, (size_t)m * k); + w->s = ei_aligned_alloc(64, ((size_t)m * kb * sizeof(ei_fp16) + 63) & ~(size_t)63); + w->c = ei_aligned_alloc(64, ((size_t)m * kb * sizeof(int16_t) + 63) & ~(size_t)63); + g2v_repack_job job = {w, blocks, kb}; + ei_parallel_for(pool, m / 8, 8, g2v_repack_groups, &job); + g2_release_source(blocks, (size_t)m * kb * sizeof(ei_block_q8_0)); + w->key = key; + w->m = m; + w->k = k; + g2v_count++; + return w; +} + +#define AVXVNNI __attribute__((target("avx2,fma,f16c,avxvnni"))) + +/* Explicitly unrolled (the build is -O2, where GCC keeps small fixed loops + * rolled and spills the accumulator arrays to the stack). The g2x_ prefix keeps + * these inside scripts/check-avx2-isolation.py's dispatched set. Each token + * carries two interleaved int32 chains (even/odd k-steps), so a 4-token tile + * keeps 8 independent vpdpbusd chains in flight (14 ymm registers). */ +#define G2V_UNROLL _Pragma("GCC unroll 16") +#define G2V_TILE(NT) \ + AVXVNNI static void g2x_v256_tile##NT(const g2v_weight *w, int g, const uint8_t *xu, const float *xs, \ + int t0, int k, float *y, int m) { \ + const int kb = k / 32; \ + const int8_t *wq = w->q + (size_t)g * kb * 256; \ + const ei_fp16 *ws = w->s + (size_t)g * kb * 8; \ + const int16_t *wc = w->c + (size_t)g * kb * 8; \ + const uint8_t *xr[NT]; \ + const float *sr[NT]; \ + __m256 acc[NT]; \ + G2V_UNROLL for (int t = 0; t < NT; t++) { \ + xr[t] = xu + (size_t)(t0 + t) * k; \ + sr[t] = xs + (size_t)(t0 + t) * kb; \ + acc[t] = _mm256_setzero_ps(); \ + } \ + for (int p = 0; p < kb; p++) { \ + __m256i even[NT], odd[NT]; \ + const __m256i corr = _mm256_slli_epi32( \ + _mm256_sub_epi32(_mm256_setzero_si256(), \ + _mm256_cvtepi16_epi32(_mm_load_si128((const __m128i *)(wc + (size_t)p * 8)))), \ + 7); /* -128 * sum(w) */ \ + G2V_UNROLL for (int t = 0; t < NT; t++) { \ + even[t] = corr; \ + odd[t] = _mm256_setzero_si256(); \ + } \ + const int8_t *wp = wq + (size_t)p * 256; \ + G2V_UNROLL for (int j = 0; j < 8; j += 2) { \ + const __m256i w0 = _mm256_load_si256((const __m256i *)(wp + j * 32)); \ + G2V_UNROLL for (int t = 0; t < NT; t++) { \ + int32_t four; \ + memcpy(&four, xr[t] + p * 32 + j * 4, 4); \ + even[t] = _mm256_dpbusd_avx_epi32(even[t], _mm256_set1_epi32(four), w0); \ + } \ + const __m256i w1 = _mm256_load_si256((const __m256i *)(wp + j * 32 + 32)); \ + G2V_UNROLL for (int t = 0; t < NT; t++) { \ + int32_t four; \ + memcpy(&four, xr[t] + p * 32 + j * 4 + 4, 4); \ + odd[t] = _mm256_dpbusd_avx_epi32(odd[t], _mm256_set1_epi32(four), w1); \ + } \ + } \ + const __m256 sw = _mm256_cvtph_ps(_mm_load_si128((const __m128i *)(ws + (size_t)p * 8))); \ + G2V_UNROLL for (int t = 0; t < NT; t++) \ + acc[t] = _mm256_fmadd_ps(_mm256_cvtepi32_ps(_mm256_add_epi32(even[t], odd[t])), \ + _mm256_mul_ps(sw, _mm256_broadcast_ss(sr[t] + p)), acc[t]);\ + } \ + G2V_UNROLL for (int t = 0; t < NT; t++) \ + _mm256_storeu_ps(y + (size_t)(t0 + t) * m + 8 * g, acc[t]); \ + } +G2V_TILE(1) +G2V_TILE(2) +G2V_TILE(3) +G2V_TILE(4) + +void g2v_tile(const void *wp, int g, const uint8_t *xu, const float *xs, int t0, int nt, int k, float *y, + int m) { + const g2v_weight *w = wp; + switch (nt) { + case 1: + g2x_v256_tile1(w, g, xu, xs, t0, k, y, m); + break; + case 2: + g2x_v256_tile2(w, g, xu, xs, t0, k, y, m); + break; + case 3: + g2x_v256_tile3(w, g, xu, xs, t0, k, y, m); + break; + default: + g2x_v256_tile4(w, g, xu, xs, t0, k, y, m); + break; + } +} +#else +int g2v_enabled(void) { return 0; } +const void *g2v_weight_of(const void *key, const ei_block_q8_0 *blocks, int m, int k, ei_thread_pool *pool) { + (void)key, (void)blocks, (void)m, (void)k, (void)pool; + return NULL; +} +void g2v_tile(const void *w, int g, const uint8_t *xu, const float *xs, int t0, int nt, int k, float *y, + int m) { + (void)w, (void)g, (void)xu, (void)xs, (void)t0, (void)nt, (void)k, (void)y, (void)m; +} +#endif diff --git a/src/gemma2_vnni256.h b/src/gemma2_vnni256.h new file mode 100644 index 0000000..ad711f0 --- /dev/null +++ b/src/gemma2_vnni256.h @@ -0,0 +1,19 @@ +/* Repacked 256-bit AVX-VNNI Q8_0 x Q8_0 GEMM (CPUs with AVX-VNNI + * but without AVX-512, e.g. Intel Alder Lake / Raptor Lake hybrid parts). */ +#ifndef EI_GEMMA2_VNNI256_H +#define EI_GEMMA2_VNNI256_H +#include "quants.h" +#include "parallel.h" +#include +#include + +/* 1 when the kernel can run here and EI_GEMMA2_CPU_VNNI256 != 0. */ +int g2v_enabled(void); +/* Repacked weights for `key` (cached; NULL if m % 8 or k % 32 or the cache is full). */ +const void *g2v_weight_of(const void *key, const ei_block_q8_0 *blocks, int m, int k, ei_thread_pool *pool); +/* y[t][8g .. 8g+7] for tokens [t0, t0 + nt), nt in 1..G2V_TILE_TOKENS; activations as + * produced by g2x_acts (u8 = q ^ 0x80 rows of k bytes, fp32 block scales). */ +enum { G2V_TILE_TOKENS = 4 }; +void g2v_tile(const void *w, int g, const uint8_t *xu, const float *xs, int t0, int nt, int k, float *y, + int m); +#endif diff --git a/src/gguf.c b/src/gguf.c index 2940fcf..f6e6a63 100644 --- a/src/gguf.c +++ b/src/gguf.c @@ -120,6 +120,37 @@ static uint64_t tensor_bytes(const ei_tensor *t, const char *namebuf) { return n * bb; } +/* Read-only file mappings currently open, so callers can tell a tensor in a + * GGUF mapping (pages re-readable from the file) from anonymous memory (where + * dropping pages would lose data). */ +#include +enum { EI_GGUF_MAPS = 16 }; +static struct { + const uint8_t *base; + size_t len; +} gguf_maps[EI_GGUF_MAPS]; +static pthread_mutex_t gguf_maps_lock = PTHREAD_MUTEX_INITIALIZER; +static void gguf_maps_set(const uint8_t *base, size_t len, const uint8_t *old) { + pthread_mutex_lock(&gguf_maps_lock); + for (int i = 0; i < EI_GGUF_MAPS; i++) + if (gguf_maps[i].base == old) { + gguf_maps[i].base = base; + gguf_maps[i].len = base ? len : 0; + break; + } + pthread_mutex_unlock(&gguf_maps_lock); +} +bool ei_gguf_in_mapping(const void *p, size_t n) { + const uint8_t *a = p; + bool found = false; + pthread_mutex_lock(&gguf_maps_lock); + for (int i = 0; i < EI_GGUF_MAPS && !found; i++) + found = gguf_maps[i].base && a >= gguf_maps[i].base && + a + n <= gguf_maps[i].base + gguf_maps[i].len; + pthread_mutex_unlock(&gguf_maps_lock); + return found; +} + void ei_gguf_open(ei_gguf *g, const char *path) { memset(g, 0, sizeof *g); @@ -146,8 +177,8 @@ void ei_gguf_open(ei_gguf *g, const char *path) { g->map = mmap(NULL, g->map_len, PROT_READ, MAP_PRIVATE, fd, 0); close(fd); if (g->map == MAP_FAILED) ei_die("%s: mmap failed", path); - #endif + gguf_maps_set(g->map, g->map_len, NULL); cur c = { g->map, g->map + g->map_len, path }; @@ -224,6 +255,7 @@ void ei_gguf_close(ei_gguf *g) { } free(g->kv); free(g->tensors); + if (g->map) gguf_maps_set(NULL, 0, g->map); #ifdef _WIN32 if (g->map) UnmapViewOfFile(g->map); #else diff --git a/src/gguf.h b/src/gguf.h index 0d97d4e..4042f30 100644 --- a/src/gguf.h +++ b/src/gguf.h @@ -61,6 +61,8 @@ typedef struct { /* Load and index a GGUF file; dies with a message on any malformation. */ void ei_gguf_open(ei_gguf *g, const char *path); void ei_gguf_close(ei_gguf *g); +/* True when [p, p + n) lies inside an open GGUF's read-only file mapping. */ +bool ei_gguf_in_mapping(const void *p, size_t n); /* Lookups (die when `required`, else return NULL/defaults). */ const ei_kv *ei_gguf_kv(const ei_gguf *g, const char *key, bool required); diff --git a/src/http_media.h b/src/http_media.h index 38eca41..6ef97fe 100644 --- a/src/http_media.h +++ b/src/http_media.h @@ -146,12 +146,20 @@ static bool media_part(const char **p, ei_media_input *m, char *err, size_t len) ei_media_frame *prepared = ei_xcalloc(frames.n, sizeof *prepared); video_decode_all(m, &frames, frame_tokens, prepared); ei_media_frames_encode(m, prepared, frames.n); + /* The "MM:SS " text before each frame matches llama.cpp (the default); + * EI_GEMMA2_SEMANTICS=official drops it, the processor's add_timestamps=False + * layout, and EI_GEMMA2_VIDEO_TIMESTAMPS=0|1 overrides either way. */ + static int stamps = -1; + if (stamps < 0) { + const char *e = getenv("EI_GEMMA2_VIDEO_TIMESTAMPS"); + stamps = e ? strcmp(e, "0") != 0 : !ei_gemma2_official_semantics(); + } for (size_t i = 0; ok && i < frames.n; ++i) { unsigned seconds = (unsigned)(i / fps); char stamp[40]; int n = snprintf(stamp, sizeof stamp, "%s%02u:%02u ", i ? " " : "", seconds / 60, seconds % 60); - ok = ei_media_text(m, stamp, (size_t)n, err, len) && + ok = (!stamps || ei_media_text(m, stamp, (size_t)n, err, len)) && ei_media_image_commit(m, &prepared[i], frames.items[i].s, frames.items[i].len, err, len); } @@ -217,6 +225,8 @@ static bool embed_structured(const server_opts *opts, const sv_string *input, fl m.cuda = opts->engine->gemma2_cuda_media ? opts->engine->gemma2_cuda_media : opts->engine->gemma2_cuda; m.lock = &inference_mutex; + if (!m.metal && !m.cuda) + m.outer_lock = &cpu_media_mutex; /* see cpu_tower_yield */ } const char *p = input->s; bool ok = false; diff --git a/src/media.c b/src/media.c index 691b816..14ab060 100644 --- a/src/media.c +++ b/src/media.c @@ -6,6 +6,7 @@ #include "gemma2_cuda.h" #endif #include "media.h" +#include "gemma2.h" #include #include #include @@ -326,6 +327,19 @@ static void cache_soft_store(const ei_media_input *m, const unsigned char key[24 ei_response_cache_insert(m->cache, (const char *)key, 24, blob, bytes); free(blob); } +/* Engine sections of a media request (outer_lock, then lock). */ +static void engine_lock(ei_media_input *m) { + if (m->outer_lock) + pthread_mutex_lock(m->outer_lock); + if (m->lock) + pthread_mutex_lock(m->lock); +} +static void engine_unlock(ei_media_input *m) { + if (m->lock) + pthread_mutex_unlock(m->lock); + if (m->outer_lock) + pthread_mutex_unlock(m->outer_lock); +} bool ei_media_image(ei_media_input *m, const char *base64, size_t len, int budget, char *err, size_t err_len) { size_t size = 0; @@ -406,8 +420,7 @@ bool ei_media_image(ei_media_input *m, const char *base64, size_t len, int budge float *soft = NULL; size_t n = 0; bool ok; - if (m->lock) - pthread_mutex_lock(m->lock); + engine_lock(m); #if defined(EI_ENABLE_METAL) if (m->metal) ok = ei_gemma2_metal_vision(m->metal, m->projector, canvas, tw, th, &soft, &n, err, err_len); else @@ -417,8 +430,7 @@ bool ei_media_image(ei_media_input *m, const char *base64, size_t len, int budge else #endif ok = ei_gemma2_vision(m->projector, m->pool, canvas, tw, th, &soft, &n, err, err_len); - if (m->lock) - pthread_mutex_unlock(m->lock); + engine_unlock(m); MEDIA_STAGE("vision", t0); free(canvas); if (ok && keyed) @@ -528,12 +540,10 @@ void ei_media_frames_encode(ei_media_input *m, ei_media_frame *frames, size_t n) float *soft = NULL; size_t total = 0; char err[160]; - if (m->lock) - pthread_mutex_lock(m->lock); + engine_lock(m); bool ok = ei_gemma2_metal_vision_frames(m->metal, m->projector, rgb, (int)batch, f->width, f->height, &soft, &total, err, sizeof err); - if (m->lock) - pthread_mutex_unlock(m->lock); + engine_unlock(m); if (ok) { size_t each = total / batch; for (size_t j = 0; j < batch; j++) { @@ -588,8 +598,7 @@ bool ei_media_image_commit(ei_media_input *m, ei_media_frame *frame, const char return true; } size_t n = 0; - if (m->lock) - pthread_mutex_lock(m->lock); + engine_lock(m); #if defined(EI_ENABLE_METAL) if (m->metal) ok = ei_gemma2_metal_vision(m->metal, m->projector, canvas, frame->width, frame->height, &soft, &n, err, err_len); else @@ -599,8 +608,7 @@ bool ei_media_image_commit(ei_media_input *m, ei_media_frame *frame, const char else #endif ok = ei_gemma2_vision(m->projector, m->pool, canvas, frame->width, frame->height, &soft, &n, err, err_len); - if (m->lock) - pthread_mutex_unlock(m->lock); + engine_unlock(m); free(canvas); if (ok && frame->keyed) cache_soft_store(m, frame->key, soft, n); @@ -618,7 +626,7 @@ static uint16_t u16(const uint8_t *p) { /* Radix-2 FFT, unnormalized, for the model's 512-point STFT. Twiddles are * tabulated once with the same expression the butterflies used to evaluate * per call (2,304 cosf/sinf pairs per frame), so results are unchanged. */ -static float fft_cos[512], fft_sin[512]; +static float fft_cos[512], fft_sin[512], hann[320]; static pthread_once_t fft_once = PTHREAD_ONCE_INIT; static void fft_tables(void) { for (int width = 2; width <= 512; width *= 2) @@ -627,6 +635,9 @@ static void fft_tables(void) { fft_cos[width / 2 + j] = cosf(theta); fft_sin[width / 2 + j] = sinf(theta); } + /* The analysis window: the same expression once instead of per frame. */ + for (int j = 0; j < 320; ++j) + hann[j] = 0.5f - 0.5f * cosf(6.283185307179586f * j / 320); } static void fft(float re[512], float im[512]) { pthread_once(&fft_once, fft_tables); @@ -655,6 +666,47 @@ static void fft(float re[512], float im[512]) { im[a] += v; } } +static bool logmel_hf(void) { + static int on = -1; + if (on < 0) { + const char *e = getenv("EI_GEMMA2_AUDIO_LOGMEL"); /* explicit override: hf | max */ + on = e ? !strcmp(e, "hf") : ei_gemma2_official_semantics(); + } + return on; +} + +/* Log-mel frames [begin, end) of one chunk; frames are + * independent, so the CPU path runs them on the engine pool. */ +typedef struct { + const float *pcm, (*filters)[257]; + const int *first, *last; + size_t count; + float *mel; + bool hf; +} mel_job; +static void mel_frames(void *opaque, int32_t begin, int32_t end) { + const mel_job *j = opaque; + for (int t = begin; t < end; ++t) { + float re[512] = {0}, im[512] = {0}, mag[257]; + for (int k = 0; k < 320; ++k) { + int pos = t * 160 + k - 160; + if (pos >= 0 && (size_t)pos < j->count) + re[k] = j->pcm[pos] * hann[k]; + } + fft(re, im); + for (int f = 0; f < 257; ++f) + mag[f] = sqrtf(re[f] * re[f] + im[f] * im[f]); + for (int bin = 0; bin < 128; ++bin) { + double sum = 0; + for (int f = j->first[bin]; f <= j->last[bin]; ++f) + sum += mag[f] * j->filters[bin][f]; + /* The default matches llama.cpp: log(max(mel, floor)). + * EI_GEMMA2_SEMANTICS=official (or EI_GEMMA2_AUDIO_LOGMEL=hf) uses the + * Transformers extractor's log(mel + floor). */ + j->mel[(size_t)t * 128 + bin] = j->hf ? (float)log(sum + 0.001) : logf(fmaxf(0.001f, (float)sum)); + } + } +} bool ei_media_audio(ei_media_input *m, const char *base64, size_t len, char *err, size_t err_len) { size_t size = 0; double t0 = media_timing() ? media_now_ms() : 0; @@ -765,32 +817,32 @@ bool ei_media_audio(ei_media_input *m, const char *base64, size_t len, char *err /* A sub-frame tail still carries signal; pad it to one analysis frame. */ if (frames < 1) frames = 1; MEDIA_STAGE("wav", t0); + /* On the CPU tower path the frames run on the engine + * pool, so the engine lock is taken here (it covers the tower anyway). */ + bool cpu_tower = !m->audio; +#if defined(EI_ENABLE_METAL) + if (m->metal) cpu_tower = false; +#endif +#if defined(EI_ENABLE_CUDA) + if (m->cuda) cpu_tower = false; +#endif + if (cpu_tower) + engine_lock(m); + pthread_once(&fft_once, fft_tables); float *mel = ei_xcalloc((size_t)frames * 128, sizeof(float)); - for (int t = 0; t < frames; ++t) { - float re[512] = {0}, im[512] = {0}, mag[257]; - for (int j = 0; j < 320; ++j) { - int pos = t * 160 + j - 160; - if (pos >= 0 && (size_t)pos < count) - re[j] = pcm[offset + pos] * (0.5f - 0.5f * cosf(6.283185307179586f * j / 320)); - } - fft(re, im); - for (int f = 0; f < 257; ++f) - mag[f] = sqrtf(re[f] * re[f] + im[f] * im[f]); - for (int bin = 0; bin < 128; ++bin) { - double sum = 0; - for (int f = first[bin]; f <= last[bin]; ++f) - sum += mag[f] * filters[bin][f]; - mel[(size_t)t * 128 + bin] = logf(fmaxf(0.001f, (float)sum)); - } - } + mel_job mj = {pcm + offset, (const float(*)[257])filters, first, last, count, mel, logmel_hf()}; + if (cpu_tower) + ei_parallel_for(m->pool, frames, 16, mel_frames, &mj); + else + mel_frames(&mj, 0, frames); MEDIA_STAGE("logmel", t0); float *soft = NULL; size_t n = 0; if (m->audio) ok = m->audio(m->audio_ctx, mel, frames, &soft, &n, err, err_len); else { - if (m->lock) - pthread_mutex_lock(m->lock); + if (!cpu_tower) + engine_lock(m); #if defined(EI_ENABLE_METAL) if (m->metal) ok = ei_gemma2_metal_audio(m->metal, m->projector, mel, frames, &soft, &n, err, err_len); else @@ -800,8 +852,7 @@ bool ei_media_audio(ei_media_input *m, const char *base64, size_t len, char *err else #endif ok = ei_gemma2_audio(m->projector, m->pool, mel, frames, &soft, &n, err, err_len); - if (m->lock) - pthread_mutex_unlock(m->lock); + engine_unlock(m); } MEDIA_STAGE("audio", t0); free(mel); @@ -834,11 +885,9 @@ bool ei_media_finish(ei_media_input *m, float out[768], char *err, size_t err_le if (m->finish) ok = m->finish(m->finish_ctx, m->inputs, m->n, out, err, err_len); else { - if (m->lock) - pthread_mutex_lock(m->lock); + engine_lock(m); ok = finish_text(m, out, err, err_len); - if (m->lock) - pthread_mutex_unlock(m->lock); + engine_unlock(m); } if (media_timing()) fprintf(stderr, "media finish %7.2f ms (%zu rows)\n", media_now_ms() - t0, m->n); diff --git a/src/media.h b/src/media.h index bf66e70..3c5e3e4 100644 --- a/src/media.h +++ b/src/media.h @@ -14,6 +14,11 @@ typedef struct { /* Optional: held only around engine calls (towers and the final text * pass), so decoding and feature extraction run concurrently. */ pthread_mutex_t *lock; + /* Optional: taken before `lock` around the same engine calls + * on the CPU backend, so one media request at a time is inside the towers + * and a tower can hand `lock` to waiting text waves between layers + * (ei_gemma2_cpu_yield) without letting another media request in. */ + pthread_mutex_t *outer_lock; /* Optional: runs the final text pass over the assembled input rows * (e.g. packed with other requests' finishes) instead of the engine call. */ bool (*finish)(void *ctx, const float *inputs, size_t n, float *out, char *err, diff --git a/src/parallel.c b/src/parallel.c index 7fc2f21..6e9b55a 100644 --- a/src/parallel.c +++ b/src/parallel.c @@ -26,6 +26,11 @@ struct ei_thread_pool { worker_context *contexts; int32_t worker_count; int32_t total_threads; + /* At most `cap` threads (caller included) join a call; + * `performance` leading threads of a hybrid-ordered pinned pool sit on + * P-cores (0 when unknown). */ + _Atomic int32_t cap; + int32_t performance; int32_t *cpus; /* pinned pools: [0] the caller, [1 + i] worker i */ pthread_mutex_t mutex; @@ -263,6 +268,7 @@ ei_thread_pool *ei_thread_pool_create(int32_t requested_threads) { ei_thread_pool *pool = calloc(1, sizeof *pool); if (!pool) ei_die("out of memory creating CPU thread pool"); pool->total_threads = total; + atomic_store(&pool->cap, total); pool->worker_count = total - 1; if (create_cpus) { pool->cpus = malloc((size_t)total * sizeof *pool->cpus); @@ -315,6 +321,84 @@ void ei_thread_pool_destroy(ei_thread_pool *pool) { free(pool); } +void ei_thread_pool_set_cap(ei_thread_pool *pool, int32_t threads) { + if (!pool) return; + if (threads < 1 || threads > pool->total_threads) threads = pool->total_threads; + atomic_store_explicit(&pool->cap, threads, memory_order_relaxed); +} + +int32_t ei_thread_pool_performance_threads(const ei_thread_pool *pool) { + return pool ? pool->performance : 0; +} + +void ei_thread_pool_set_performance_threads(ei_thread_pool *pool, int32_t threads) { + if (pool) pool->performance = threads < 0 ? 0 : threads > pool->total_threads ? pool->total_threads : threads; +} + +#if defined(__linux__) +static int read_cpulist(const char *path, unsigned char *mask, int max) { + FILE *f = fopen(path, "r"); + if (!f) return 0; + char buf[1024]; + int n = 0; + if (fgets(buf, sizeof buf, f)) { + for (char *p = buf; *p && *p != '\n';) { + long a = strtol(p, &p, 10), b = a; + if (*p == '-') b = strtol(p + 1, &p, 10); + for (long c = a; c <= b && c < max; c++) mask[c] = 1, n++; + if (*p == ',') p++; + else if (*p && *p != '\n') break; + } + } + fclose(f); + return n; +} +#endif + +int32_t ei_cpu_hybrid_order(int32_t *cpus, int32_t max, int32_t *performance) { + *performance = 0; +#if defined(__linux__) + enum { M = 4096 }; + static unsigned char pcore[M], ecore[M]; + memset(pcore, 0, sizeof pcore), memset(ecore, 0, sizeof ecore); + if (!read_cpulist("/sys/devices/cpu_core/cpus", pcore, M) || + !read_cpulist("/sys/devices/cpu_atom/cpus", ecore, M)) + return 0; /* not a hybrid CPU (or not reported): keep upstream's placement */ + cpu_set_t allowed; + if (sched_getaffinity(0, sizeof allowed, &allowed) != 0) return 0; + int32_t n = ei_cpu_count(), out = 0; + int32_t *primary_core = calloc((size_t)n, sizeof *primary_core); + if (!primary_core) return 0; + /* Pass 0: P-core primary threads, 1: E-cores, 2: P-core SMT siblings. */ + for (int pass = 0; pass < 3; pass++) { + int32_t seen = 0; + for (int32_t c = 0; c < n && c < M; c++) { + if (!CPU_ISSET(c, &allowed)) continue; + int core = read_sys_int("/sys/devices/system/cpu/cpu%d/topology/core_id", c); + bool sibling = false; + for (int32_t j = 0; j < seen; j++) + if (primary_core[j] == core) sibling = true; + if (pcore[c] && !sibling && pass == 0 && out < max) { + cpus[out++] = c; + primary_core[seen++] = core; + (*performance)++; + } else if (pcore[c] && !sibling) { + primary_core[seen++] = core; + } else if (pcore[c] && sibling && pass == 2 && out < max) { + cpus[out++] = c; + } else if (ecore[c] && pass == 1 && out < max) { + cpus[out++] = c; + } + } + } + free(primary_core); + return out; +#else + (void)cpus, (void)max; + return 0; +#endif +} + int32_t ei_thread_pool_threads(const ei_thread_pool *pool) { return pool ? pool->total_threads : 1; } @@ -328,9 +412,11 @@ void ei_parallel_for(ei_thread_pool *pool, int32_t count, int32_t grain, } if (grain < 1) grain = 1; + const int32_t cap = atomic_load_explicit(&pool->cap, memory_order_relaxed); if (pool->spin_ns) { int32_t jobs = (count + grain - 1) / grain, helpers = jobs - 1; if (helpers > pool->worker_count) helpers = pool->worker_count; + if (helpers > cap - 1) helpers = cap - 1 > 0 ? cap - 1 : 0; pool->count = count; pool->grain = grain; pool->function = function; @@ -360,6 +446,7 @@ void ei_parallel_for(ei_thread_pool *pool, int32_t count, int32_t grain, pool->context = context; int32_t jobs = (count + grain - 1) / grain; pool->active_workers = jobs - 1; + if (pool->active_workers > cap - 1) pool->active_workers = cap - 1 > 0 ? cap - 1 : 0; if (pool->active_workers > pool->worker_count) { pool->active_workers = pool->worker_count; } diff --git a/src/parallel.h b/src/parallel.h index 3828218..39a9fdf 100644 --- a/src/parallel.h +++ b/src/parallel.h @@ -25,6 +25,17 @@ void ei_thread_pool_unbind_caller(const ei_cpu_binding *saved); /* Logical CPUs in placement order: one hardware thread per core first, grouped * by shared L3, then SMT siblings (Linux sysfs; identity elsewhere). */ int32_t ei_cpu_placement(int32_t *cpus, int32_t max); +/* Hybrid CPUs. Allowed CPUs (the process affinity mask) in the + * order P-core primary threads, E-cores, P-core SMT siblings; *performance = + * number of P-core primaries. Returns 0 when the CPU does not report hybrid + * core types (/sys/devices/cpu_core, cpu_atom). */ +int32_t ei_cpu_hybrid_order(int32_t *cpus, int32_t max, int32_t *performance); +/* At most `threads` threads (caller included) join subsequent calls; values + * outside 1..threads() restore the full pool. Set between calls only. */ +void ei_thread_pool_set_cap(ei_thread_pool *pool, int32_t threads); +/* Leading pool threads pinned to P-cores (0 = unknown / not hybrid). */ +int32_t ei_thread_pool_performance_threads(const ei_thread_pool *pool); +void ei_thread_pool_set_performance_threads(ei_thread_pool *pool, int32_t threads); void ei_thread_pool_destroy(ei_thread_pool *pool); int32_t ei_thread_pool_threads(const ei_thread_pool *pool); /* Online logical CPUs (at least 1). */ diff --git a/src/server.c b/src/server.c index 8c2ae30..fcb1284 100644 --- a/src/server.c +++ b/src/server.c @@ -65,6 +65,50 @@ static int ei_setsockopt(ei_socket fd, int level, int name, } static pthread_mutex_t inference_mutex = PTHREAD_MUTEX_INITIALIZER; +/* On the CPU backend a media tower holds inference_mutex for seconds (a + * 1120-token image), and every text wave queued behind it waited for the + * whole tower. Media requests also take cpu_media_mutex (before + * inference_mutex, so one media request at a time is inside the towers), and + * a tower calls cpu_tower_yield() between layers: when text waves are + * waiting it releases inference_mutex until one of them holds it, then waits + * its turn again. Text waves never take cpu_media_mutex or a tower lock, so + * the handoff cannot deadlock. EI_GEMMA2_CPU_YIELD=0 restores the plain lock. */ +static pthread_mutex_t cpu_media_mutex = PTHREAD_MUTEX_INITIALIZER; +static atomic_int text_waiters; /* text waves blocked on inference_mutex */ +static pthread_mutex_t handoff_mutex = PTHREAD_MUTEX_INITIALIZER; +static pthread_cond_t handoff_cond = PTHREAD_COND_INITIALIZER; +static uint64_t handoffs; /* text waves that acquired inference_mutex (handoff_mutex) */ + +static void text_engine_lock(void) { + atomic_fetch_add_explicit(&text_waiters, 1, memory_order_relaxed); + pthread_mutex_lock(&inference_mutex); + atomic_fetch_sub_explicit(&text_waiters, 1, memory_order_relaxed); + pthread_mutex_lock(&handoff_mutex); + handoffs++; + pthread_cond_broadcast(&handoff_cond); + pthread_mutex_unlock(&handoff_mutex); +} + +/* Called by a CPU media tower that holds inference_mutex (and cpu_media_mutex). */ +static void cpu_tower_yield(void) { + if (!atomic_load_explicit(&text_waiters, memory_order_relaxed)) + return; + pthread_mutex_lock(&handoff_mutex); + const uint64_t seen = handoffs; + pthread_mutex_unlock(&inference_mutex); + struct timespec deadline; + clock_gettime(CLOCK_REALTIME, &deadline); + deadline.tv_nsec += 50 * 1000000L; /* bounded: a waiter that never shows up */ + if (deadline.tv_nsec >= 1000000000L) { + deadline.tv_sec++; + deadline.tv_nsec -= 1000000000L; + } + while (handoffs == seen && atomic_load_explicit(&text_waiters, memory_order_relaxed)) + if (pthread_cond_timedwait(&handoff_cond, &handoff_mutex, &deadline) == ETIMEDOUT) + break; + pthread_mutex_unlock(&handoff_mutex); + pthread_mutex_lock(&inference_mutex); +} // GPU media contexts (vision/audio towers and the mixed-input text forward): // each slot is a clone sharing the device, pipelines and weights with its own // queue, workspace and lock, so media requests overlap on the GPU instead of @@ -1955,7 +1999,7 @@ static bool execute_engine_batch(void *opaque, const int32_t *ids, if (!engine->model.embeddinggemma2 || media_slot_count) return ei_engine_embed_tokens_batch(engine, ids, offsets, batch_size, out, err, err_len); // CPU backend: media towers share the engine's thread pool with text waves. - pthread_mutex_lock(&inference_mutex); + text_engine_lock(); bool ok = ei_engine_embed_tokens_batch(opaque, ids, offsets, batch_size, out, err, err_len); pthread_mutex_unlock(&inference_mutex); return ok; @@ -1970,6 +2014,43 @@ static bool execute_engine_batch_slot(void *opaque, size_t slot, const int32_t * err_len); } +/* EI_GEMMA2_CPU_PREPACK=1 runs a one-token text, a 48x48 image and a short + * silent clip through the CPU graphs before the server listens, so the + * one-time weight repack (0.3-0.4 s per tower) is not paid by the first + * request of each modality. Off by default: the repacked copies of towers a + * deployment never uses would stay resident. */ +static void cpu_prepack(ei_engine *e, const ei_gguf *projector) { + const char *v = getenv("EI_GEMMA2_CPU_PREPACK"); + if (!v || !strcmp(v, "0") || !e->model.embeddinggemma2 || strcmp(ei_engine_backend(e), "cpu")) + return; + struct timespec a, b; + clock_gettime(CLOCK_MONOTONIC, &a); + char err[256]; + float out[768], *soft = NULL; + size_t n = 0, offsets[2] = {0, 1}; + int32_t id = e->model.bos_id; + bool ok = ei_gemma2_embed_batch_fast(&e->model.gguf, e->thread_pool, &id, NULL, offsets, 1, out, err, + sizeof err); + if (ok && projector) { + uint8_t rgb[48 * 48 * 3]; + memset(rgb, 128, sizeof rgb); + ok = ei_gemma2_vision(projector, e->thread_pool, rgb, 48, 48, &soft, &n, err, sizeof err); + free(soft), soft = NULL; + } + if (ok && projector) { + float mel[16 * 128]; + for (size_t i = 0; i < 16 * 128; i++) + mel[i] = -6.9f; /* log(0.001): silence */ + ok = ei_gemma2_audio(projector, e->thread_pool, mel, 16, &soft, &n, err, sizeof err); + free(soft); + } + clock_gettime(CLOCK_MONOTONIC, &b); + if (!ok) + ei_die("EI_GEMMA2_CPU_PREPACK: %s", err); + fprintf(stderr, "prepacked CPU weights (text%s) in %.0f ms\n", projector ? ", vision, audio" : "", + (b.tv_sec - a.tv_sec) * 1e3 + (b.tv_nsec - a.tv_nsec) * 1e-6); +} + int main(int argc, char **argv) { #ifdef _WIN32 WSADATA wsa; @@ -2128,6 +2209,13 @@ int main(int argc, char **argv) { } #endif } + if (second && ei_gemma2_cpu_fast()) + cpu_prepack(&engine, opts.projector); + if (second && cpu_backend && !media_slot_count) { + const char *yield = getenv("EI_GEMMA2_CPU_YIELD"); + if (!(yield && !strcmp(yield, "0"))) + ei_gemma2_set_cpu_yield(cpu_tower_yield); + } // Each keep-alive connection holds an HTTP worker, so workers cap the // requests in flight; give every executor enough to fill its waves. if (!opts.workers_explicit && opts.workers < 32 * engine.executor_count) { diff --git a/src/test_gemma2_attn.c b/src/test_gemma2_attn.c new file mode 100644 index 0000000..b695700 --- /dev/null +++ b/src/test_gemma2_attn.c @@ -0,0 +1,179 @@ +/* g2m_attn (AVX2 tiled online softmax) against its scalar + * reference on random data: head widths 64/128/256/512, GQA strides, + * sliding windows and packed-sequence bounds, partial query/key blocks. + * Both are compared with a float64 ground truth (max |diff| in a row / row RMS); + * exits 1 if the tiled kernel is more than 2x less accurate than plain scalar + * fp32 or above 1e-3. */ +#if defined(__linux__) +#define _POSIX_C_SOURCE 200809L +#define _DEFAULT_SOURCE /* madvise */ +#endif +#include "gemma2_media_x86.h" +#include "common.h" +#include "parallel.h" +#include +#include +#if defined(__linux__) +#include +#endif +#include +#include +#include + +static uint64_t rng = 0x243F6A8885A308D3ull; +static float frand(void) { + rng ^= rng << 13, rng ^= rng >> 7, rng ^= rng << 17; + return (float)((rng >> 11) * (1.0 / 9007199254740992.0)) * 2.0f - 1.0f; +} + +/* max over (token, head) rows of max|a - ref| / rms(ref) */ +static double row_err(const float *a, const double *ref, int n, int heads, int d) { + double worst = 0; + for (size_t r = 0; r < (size_t)n * heads; r++) { + double ss = 0, md = 0; + for (int z = 0; z < d; z++) { + double x = a[r * d + z], y = ref[r * d + z]; + ss += y * y; + md = x != x ? INFINITY : fabs(x - y) > md ? fabs(x - y) : md; + } + double e = md / sqrt(ss / d + 1e-30); + if (!(e <= worst)) + worst = e; + } + return worst; +} +static void run(int d, int n, int heads, int kvh, int window, int split, double *err_avx2, double *err_scalar) { + size_t qs = (size_t)heads * d, ks = (size_t)kvh * d; + float *q = ei_xmalloc(n * qs * 4), *k = ei_xmalloc(n * ks * 4), *v = ei_xmalloc(n * ks * 4); + float *o1 = ei_xcalloc(n * qs, 4), *o2 = ei_xcalloc(n * qs, 4); + double *exact = ei_xcalloc(n * qs, sizeof(double)), *sc = ei_xcalloc(n, sizeof(double)); + /* Unit-scale rows as after the model's q/k/v RMS norms; scores reach ~d. */ + for (size_t i = 0; i < n * qs; i++) + q[i] = frand() * 1.7f; + for (size_t i = 0; i < n * ks; i++) + k[i] = frand() * 1.7f, v[i] = frand() * 1.7f; + int qb = g2m_attn_qblock(d); + int32_t lo[G2M_ATTN_QB_MAX], hi[G2M_ATTN_QB_MAX]; + for (int hd = 0; hd < heads; hd++) + for (int q0 = 0; q0 < n; q0 += qb) { + int nq = n - q0 < qb ? n - q0 : qb, klo = n, khi = 0; + for (int i = 0; i < nq; i++) { + int t = q0 + i, s0 = t < split ? 0 : split, s1 = t < split ? split : n; + lo[i] = window && t - window > s0 ? t - window : s0; + hi[i] = window && t + window + 1 < s1 ? t + window + 1 : s1; + klo = lo[i] < klo ? lo[i] : klo; + khi = hi[i] > khi ? hi[i] : khi; + /* double-precision ground truth */ + const float *qi = q + (size_t)t * qs + (size_t)hd * d; + int kh = hd / (heads / kvh); + double mx = -INFINITY, total = 0; + for (int j = lo[i]; j < hi[i]; j++) { + double dot = 0; + for (int z = 0; z < d; z++) + dot += (double)qi[z] * k[(size_t)j * ks + (size_t)kh * d + z]; + sc[j] = dot; + mx = dot > mx ? dot : mx; + } + double *ex = exact + (size_t)t * qs + (size_t)hd * d; + for (int j = lo[i]; j < hi[i]; j++) { + double p = exp(sc[j] - mx); + total += p; + for (int z = 0; z < d; z++) + ex[z] += p * v[(size_t)j * ks + (size_t)kh * d + z]; + } + for (int z = 0; z < d; z++) + ex[z] /= total; + } + g2m_attn_args a = {.q = q, .k = k, .v = v, .qstride = qs, .kvstride = ks, .ostride = qs, + .qcol = (size_t)hd * d, .kvcol = (size_t)(hd / (heads / kvh)) * d, + .ocol = (size_t)hd * d, .d = d, .q0 = q0, .nq = nq, .lo = lo, .hi = hi, + .klo = klo, .khi = khi, .scale = 1.0f}; + a.out = o1; + g2m_attn(&a); + a.out = o2; + g2m_attn_reference(&a); + } + *err_avx2 = row_err(o1, exact, n, heads, d); + *err_scalar = row_err(o2, exact, n, heads, d); + free(q), free(k), free(v), free(o1), free(o2), free(exact), free(sc); +} + +/* bench N D HEADS: all heads x query blocks of full attention, pool of EI_THREADS. */ +typedef struct { + const float *q, *k, *v; + float *o; + int n, d, heads, blocks, headmajor; +} bench_job; +static void bench_tasks(void *opaque, int32_t begin, int32_t end) { + bench_job *j = opaque; + int qb = g2m_attn_qblock(j->d); + for (int task = begin; task < end; task++) { + int hd = task / j->blocks, q0 = task % j->blocks * qb, nq = j->n - q0 < qb ? j->n - q0 : qb; + size_t st = (size_t)j->heads * j->d, col = (size_t)hd * j->d, base = 0; + if (j->headmajor) /* [heads][n][d]: each head's rows contiguous */ + st = j->d, col = 0, base = (size_t)hd * j->n * j->d; + g2m_attn_args a = {.q = j->q + base, .k = j->k + base, .v = j->v + base, .out = j->o, .qstride = st, + .kvstride = st, .ostride = (size_t)j->heads * j->d, .qcol = col, .kvcol = col, + .ocol = (size_t)hd * j->d, .d = j->d, .q0 = q0, .nq = nq, .klo = 0, .khi = j->n, .scale = 1.0f}; + g2m_attn(&a); + } +} +static int bench(int n, int d, int heads) { + size_t st = (size_t)heads * d; + /* EI_ATTN_BENCH_HEADMAJOR=1: Q/K/V as [heads][n][d]; EI_ATTN_BENCH_THP=1: 2 MB-aligned, MADV_HUGEPAGE. */ + const char *hm = getenv("EI_ATTN_BENCH_HEADMAJOR"), *thp = getenv("EI_ATTN_BENCH_THP"); + float *buf[4]; + size_t bytes = ((size_t)n * st * 4 + (2u << 20) - 1) & ~(size_t)((2u << 20) - 1); + for (int i = 0; i < 4; i++) { + buf[i] = ei_aligned_alloc(2u << 20, bytes); +#if defined(__linux__) + if (thp && atoi(thp)) + madvise(buf[i], bytes, MADV_HUGEPAGE); +#else + (void)thp; +#endif + } + float *q = buf[0], *k = buf[1], *v = buf[2], *o = buf[3]; + for (size_t i = 0; i < n * st; i++) + q[i] = frand() * 0.3f, k[i] = frand() * 0.3f, v[i] = frand(); + ei_thread_pool *pool = ei_thread_pool_create_spin(0, 50); + bench_job j = {q, k, v, o, n, d, heads, (n + g2m_attn_qblock(d) - 1) / g2m_attn_qblock(d), hm && atoi(hm)}; + double best = 1e30; + for (int r = 0; r < 5; r++) { + struct timespec a, b; + clock_gettime(CLOCK_MONOTONIC, &a); + ei_parallel_for(pool, heads * j.blocks, 1, bench_tasks, &j); + clock_gettime(CLOCK_MONOTONIC, &b); + double ms = (b.tv_sec - a.tv_sec) * 1e3 + (b.tv_nsec - a.tv_nsec) * 1e-6; + best = ms < best ? ms : best; + } + double macs = 2.0 * n * (double)n * d * heads; + printf("attention n=%d d=%d heads=%d threads=%d: best %.2f ms, %.1f GMAC/s\n", n, d, heads, + ei_thread_pool_threads(pool), best, macs / (best * 1e6)); + ei_thread_pool_destroy(pool); + ei_aligned_free(q), ei_aligned_free(k), ei_aligned_free(v), ei_aligned_free(o); + return 0; +} + +int main(int argc, char **argv) { + if (argc == 5 && !strcmp(argv[1], "bench")) + return bench(atoi(argv[2]), atoi(argv[3]), atoi(argv[4])); + struct { int d, n, heads, kvh, window, split; } cases[] = { + {64, 630, 12, 12, 0, 630}, /* vision i70 */ + {64, 77, 12, 12, 0, 77}, /* partial blocks */ + {128, 300, 8, 8, 11, 300}, /* audio-like width */ + {256, 1100, 4, 2, 512, 1100}, /* text sliding layer, GQA 4:2 */ + {256, 700, 4, 2, 512, 333}, /* two packed sequences */ + {512, 900, 4, 1, 0, 900}, /* text full layer, MQA */ + {512, 300, 4, 1, 0, 37}, /* packed, short first sequence */ + }; + bool ok = true; + for (size_t c = 0; c < sizeof cases / sizeof *cases; c++) { + double ea, es; + run(cases[c].d, cases[c].n, cases[c].heads, cases[c].kvh, cases[c].window, cases[c].split, &ea, &es); + printf("d=%3d n=%5d heads=%2d kv=%2d window=%3d split=%4d err/rms vs fp64: tiled %.2e scalar fp32 %.2e\n", + cases[c].d, cases[c].n, cases[c].heads, cases[c].kvh, cases[c].window, cases[c].split, ea, es); + ok &= ea <= 2 * es + 1e-6 && ea < 1e-3; + } + return ok ? 0 : 1; +} diff --git a/src/test_gemma2_media.c b/src/test_gemma2_media.c new file mode 100644 index 0000000..bc6d80e --- /dev/null +++ b/src/test_gemma2_media.c @@ -0,0 +1,217 @@ +/* Parity and timing of the fast CPU media towers against the + * unchanged reference graph (same process, same pool, same inputs). + * + * test_gemma2_media text.gguf mmproj.gguf vision A.png [B.png ...] + * test_gemma2_media text.gguf mmproj.gguf vision-synthetic + * + * Images must already be 48-pixel aligned (the fixtures are). Prints one JSON + * line per input: soft-token count, flattened and minimum per-token cosine, + * max |diff|, reference and fast tower milliseconds (median of + * EI_MEDIA_TEST_REPEAT fast runs). EI_MEDIA_TEST_REF=0 skips the reference. + * Exit status 1 if a gate in check() fails. */ +#if defined(__linux__) +#define _POSIX_C_SOURCE 200809L +#endif +#include "engine.h" +#include "gemma2.h" +#include +#include +#include +#include +#include +#define STBI_NO_STDIO +#include "vendor/stb_image.h" + +static double now_ms(void) { + struct timespec ts; + clock_gettime(CLOCK_MONOTONIC, &ts); + return ts.tv_sec * 1e3 + ts.tv_nsec * 1e-6; +} +static uint8_t *read_file(const char *path, size_t *size) { + FILE *f = fopen(path, "rb"); + if (!f) + ei_die("cannot open %s", path); + fseek(f, 0, SEEK_END); + long n = ftell(f); + fseek(f, 0, SEEK_SET); + uint8_t *data = ei_xmalloc((size_t)n); + if (fread(data, 1, (size_t)n, f) != (size_t)n) + ei_die("cannot read %s", path); + fclose(f); + *size = (size_t)n; + return data; +} +static int cmp_double(const void *a, const void *b) { + double x = *(const double *)a, y = *(const double *)b; + return x < y ? -1 : x > y; +} +static int env_int(const char *name, int fallback) { + const char *v = getenv(name); + return v ? atoi(v) : fallback; +} + +/* Gate: flattened cosine >= EI_MEDIA_TEST_MIN_COS (default 0.999) and every token + * >= EI_MEDIA_TEST_MIN_TOKEN_COS (default 0.95). Both graphs quantize activations + * to Q8_0, so they differ by rounding noise; against the HF FP32 tower on real + * inputs each graph's worst token is ~0.95 (measured in the CPU-Max fork). */ +static bool check(const char *label, const float *ref, const float *fast, size_t tokens, double ref_ms, + double fast_ms, double min_allowed) { + double min_token_allowed = + getenv("EI_MEDIA_TEST_MIN_TOKEN_COS") ? atof(getenv("EI_MEDIA_TEST_MIN_TOKEN_COS")) : 0.95; + double dot = 0, rr = 0, ff = 0, min_tok = 1, max_diff = 0; + for (size_t t = 0; t < tokens; t++) { + double d = 0, a = 0, b = 0; + for (int z = 0; z < 512; z++) { + double x = ref[t * 512 + z], y = fast[t * 512 + z]; + d += x * y, a += x * x, b += y * y; + if (fabs(x - y) > max_diff) + max_diff = fabs(x - y); + } + dot += d, rr += a, ff += b; + double c = d / sqrt(a * b); + if (c < min_tok) + min_tok = c; + } + double cos = dot / sqrt(rr * ff); + printf("{\"input\":\"%s\",\"tokens\":%zu,\"cosine\":%.9f,\"min_token_cosine\":%.9f,\"max_abs_diff\":%.6g," + "\"ref_ms\":%.2f,\"fast_ms\":%.2f,\"speedup\":%.2f}\n", + label, tokens, cos, min_tok, max_diff, ref_ms, fast_ms, ref_ms > 0 ? ref_ms / fast_ms : 0.0); + fflush(stdout); + return cos >= min_allowed && min_tok >= min_token_allowed && isfinite(cos); +} + +typedef struct { + const uint8_t *rgb; + int w, h; + const float *mel; + int frames; +} media_input; +static bool tower(const ei_gguf *proj, ei_thread_pool *pool, const media_input *in, bool fast, float **out, + size_t *n, char *err, size_t len) { + if (in->mel) + return fast ? ei_gemma2_audio_fast(proj, pool, in->mel, in->frames, out, n, err, len) + : ei_gemma2_audio_ref(proj, pool, in->mel, in->frames, out, n, err, len); + return fast ? ei_gemma2_vision_fast(proj, pool, in->rgb, in->w, in->h, out, n, err, len) + : ei_gemma2_vision_ref(proj, pool, in->rgb, in->w, in->h, out, n, err, len); +} +static bool run_vision(ei_engine *e, const ei_gguf *proj, const char *label, const media_input *in) { + char err[256]; + int repeat = env_int("EI_MEDIA_TEST_REPEAT", 3); + bool do_ref = env_int("EI_MEDIA_TEST_REF", 1) != 0; + double min_allowed = getenv("EI_MEDIA_TEST_MIN_COS") ? atof(getenv("EI_MEDIA_TEST_MIN_COS")) : 0.999; + float *fast = NULL, *ref = NULL; + size_t nf = 0, nr = 0; + double times[32]; + if (repeat > 32) + repeat = 32; + for (int r = 0; r < repeat; r++) { + free(fast); + double t0 = now_ms(); + if (!tower(proj, e->thread_pool, in, true, &fast, &nf, err, sizeof err)) + ei_die("%s", err); + times[r] = now_ms() - t0; + } + qsort(times, (size_t)repeat, sizeof *times, cmp_double); + double ref_ms = 0; + if (do_ref) { + double t0 = now_ms(); + if (!tower(proj, e->thread_pool, in, false, &ref, &nr, err, sizeof err)) + ei_die("%s", err); + ref_ms = now_ms() - t0; + if (nr != nf) + ei_die("token count mismatch %zu vs %zu", nr, nf); + } + const char *dump = getenv("EI_MEDIA_TEST_DUMP"); /* prefix: .fast.f32 / .ref.f32 */ + if (dump) { + char path[512]; + for (int which = 0; which < 2; which++) { + const float *src = which ? ref : fast; + if (!src) + continue; + snprintf(path, sizeof path, "%s.%s.f32", dump, which ? "ref" : "fast"); + FILE *f = fopen(path, "wb"); + if (!f || fwrite(src, sizeof(float), nf * 512, f) != nf * 512) + ei_die("cannot write %s", path); + fclose(f); + } + } + bool ok = true; + if (do_ref) + ok = check(label, ref, fast, nf, ref_ms, times[repeat / 2], min_allowed); + else + printf("{\"input\":\"%s\",\"tokens\":%zu,\"fast_ms\":%.2f}\n", label, nf, times[repeat / 2]); + for (size_t i = 0; i < nf * 512; i++) + if (!isfinite(fast[i])) + ok = false; + free(fast); + free(ref); + return ok; +} + +int main(int argc, char **argv) { + if (argc < 4) + ei_die("usage: %s text.gguf mmproj.gguf vision IMAGE.png... | vision-synthetic | audio MEL.f32... | audio-synthetic", argv[0]); + ei_engine e; + ei_engine_load(&e, argv[1]); + ei_gguf proj; + ei_gguf_open(&proj, argv[2]); + bool ok = true; + if (!strcmp(argv[3], "vision-synthetic")) { + /* Gradients and a checkerboard at three aspect ratios / budgets. */ + const int sizes[][2] = {{96, 96}, {480, 336}, {144, 432}}; + for (int s = 0; s < 3; s++) { + int w = sizes[s][0], h = sizes[s][1]; + uint8_t *rgb = ei_xmalloc((size_t)w * h * 3); + for (int y = 0; y < h; y++) + for (int x = 0; x < w; x++) { + uint8_t *p = rgb + ((size_t)y * w + x) * 3; + p[0] = (uint8_t)(x * 255 / w); + p[1] = (uint8_t)(y * 255 / h); + p[2] = ((x / 8 + y / 8) & 1) ? 230 : 20; + } + char label[64]; + snprintf(label, sizeof label, "synthetic-%dx%d", w, h); + media_input in = {.rgb = rgb, .w = w, .h = h}; + ok &= run_vision(&e, &proj, label, &in); + free(rgb); + } + } else if (!strcmp(argv[3], "vision")) { + for (int i = 4; i < argc; i++) { + size_t size; + uint8_t *data = read_file(argv[i], &size); + int w, h, c; + uint8_t *rgb = stbi_load_from_memory(data, (int)size, &w, &h, &c, 3); + free(data); + if (!rgb) + ei_die("cannot decode %s", argv[i]); + media_input in = {.rgb = rgb, .w = w, .h = h}; + ok &= run_vision(&e, &proj, argv[i], &in); + stbi_image_free(rgb); + } + } else if (!strcmp(argv[3], "audio")) { + /* Raw float32 log-mel features, [frames][128] (e.g. from the HF extractor). */ + for (int i = 4; i < argc; i++) { + size_t size; + float *mel = (float *)read_file(argv[i], &size); + media_input in = {.mel = mel, .frames = (int)(size / (128 * sizeof(float)))}; + ok &= run_vision(&e, &proj, argv[i], &in); + free(mel); + } + } else if (!strcmp(argv[3], "audio-synthetic")) { + const int lens[] = {96, 499, 1499}; + for (int s = 0; s < 3; s++) { + float *mel = ei_xmalloc((size_t)lens[s] * 128 * sizeof(float)); + for (int i = 0; i < lens[s] * 128; i++) + mel[i] = sinf(i * 0.017f) * 2.0f - 4.0f; + char label[64]; + snprintf(label, sizeof label, "synthetic-mel-%d", lens[s]); + media_input in = {.mel = mel, .frames = lens[s]}; + ok &= run_vision(&e, &proj, label, &in); + free(mel); + } + } else + ei_die("unknown mode %s", argv[3]); + ei_gguf_close(&proj); + ei_engine_free(&e); + return ok ? 0 : 1; +} diff --git a/src/test_gemma2_vnni256.c b/src/test_gemma2_vnni256.c new file mode 100644 index 0000000..a6440d5 --- /dev/null +++ b/src/test_gemma2_vnni256.c @@ -0,0 +1,151 @@ +/* Exactness and throughput of the Q8_0 GEMM routes. + * + * test_gemma2_vnni256 [NxMxK ...] (default: the vision/audio/text shapes) + * fNxMxK / FNxMxK: the float-weight GEMM with BF16 / F32 weights (EI_GEMMA2_CPU_SGEMM=0: the row kernel) + * + * Random Q8_0 weights and activations; the result of ei_gemma2_fast_gemm_q8 + * (whatever route the dispatch picks: set EI_GEMMA2_CPU_VNNI256=0 for the + * generic tile) is compared with a double-precision scalar Q8_0 dot product. + * Prints max relative error and GMAC/s (median of 5) per shape; exits 1 if a + * relative error exceeds 1e-5. EI_THREADS sets the pool width. */ +#if defined(__linux__) +#define _POSIX_C_SOURCE 200809L +#endif +#include "gemma2.h" +#include "gemma2_vnni256.h" +#include "parallel.h" +#include +#include +#include +#include +#include + +static double now_ms(void) { + struct timespec ts; + clock_gettime(CLOCK_MONOTONIC, &ts); + return ts.tv_sec * 1e3 + ts.tv_nsec * 1e-6; +} +static uint64_t rng = 0x9E3779B97F4A7C15ull; +static float frand(void) { + rng ^= rng << 13, rng ^= rng >> 7, rng ^= rng << 17; + return (float)((rng >> 11) * (1.0 / 9007199254740992.0)) * 2.0f - 1.0f; +} +static int cmp_double(const void *a, const void *b) { + double x = *(const double *)a, y = *(const double *)b; + return x < y ? -1 : x > y; +} + +int main(int argc, char **argv) { + const char *defaults[] = {"2520x768x768", "2520x3072x768", "2520x768x3072", "750x4096x1024", + "750x1024x4096", "512x2048x512", "512x512x2048", "16x2048x512"}; + const char **cases = argc > 1 ? (const char **)argv + 1 : defaults; + int n_cases = argc > 1 ? argc - 1 : (int)(sizeof defaults / sizeof *defaults); + ei_thread_pool *pool = ei_thread_pool_create_spin(0, 50); + ei_gemma2_fast *ctx = ei_gemma2_fast_create(); + printf("threads=%d vnni256=%d\n", ei_thread_pool_threads(pool), g2v_enabled()); + bool ok = true; + for (int c = 0; c < n_cases; c++) { + int n, m, k; + if (cases[c][0] == 'f' || cases[c][0] == 'F') { /* float-weight GEMM: f = BF16, F = F32 weights */ + bool bf16 = cases[c][0] == 'f'; + if (sscanf(cases[c] + 1, "%dx%dx%d", &n, &m, &k) != 3) + ei_die("bad case %s", cases[c]); + float *wf = ei_xmalloc((size_t)m * k * sizeof(float)), *x = ei_xmalloc((size_t)n * k * sizeof(float)); + uint16_t *wb = ei_xmalloc((size_t)m * k * sizeof(uint16_t)); + for (size_t i = 0; i < (size_t)m * k; i++) { + float f = frand() * 0.1f; + uint32_t bits; + memcpy(&bits, &f, 4); + wb[i] = (uint16_t)(bits >> 16); + bits = (uint32_t)wb[i] << 16; + memcpy(&wf[i], &bits, 4); /* the exact value the BF16 tensor holds */ + } + for (size_t i = 0; i < (size_t)n * k; i++) + x[i] = frand() * 3.0f; + ei_tensor *wt = ei_xmalloc(sizeof *wt); + *wt = (ei_tensor){.n_dims = 2, .ne = {(uint64_t)k, (uint64_t)m, 1, 1}, + .type = bf16 ? EI_T_BF16 : EI_T_F32, .data = bf16 ? (void *)wb : (void *)wf}; + float *yf = ei_xmalloc((size_t)n * m * sizeof(float)); + double times[5]; + for (int rep = 0; rep < 6; rep++) { + double t0 = now_ms(); + ei_gemma2_fast_gemm_f(pool, wt, x, n, yf); + if (rep) + times[rep - 1] = now_ms() - t0; + } + qsort(times, 5, sizeof *times, cmp_double); + double max_rel = 0; + for (int t = 0; t < n; t += n > 64 ? 29 : 1) + for (int r = 0; r < m; r += m > 256 ? 11 : 1) { + double ref = 0, mag = 0; + for (int z = 0; z < k; z++) + ref += (double)wf[(size_t)r * k + z] * x[(size_t)t * k + z], + mag += fabs((double)wf[(size_t)r * k + z] * x[(size_t)t * k + z]); + double rel = fabs(yf[(size_t)t * m + r] - ref) / (mag > 0 ? mag : 1); + max_rel = rel > max_rel ? rel : max_rel; + } + printf("%-16s median %8.3f ms %8.1f GMAC/s max_rel_err %.2e (float weights, %s)\n", cases[c], times[2], + (double)n * m * k / (times[2] * 1e6), max_rel, bf16 ? "BF16" : "F32"); + if (!(max_rel < 1e-5)) + ok = false; + free(x), free(yf); + continue; + } + if (sscanf(cases[c], "%dx%dx%d", &n, &m, &k) != 3 || k % 32) + ei_die("bad case %s (NxMxK, K multiple of 32)", cases[c]); + ei_block_q8_0 *wq = ei_xmalloc((size_t)m * k / 32 * sizeof *wq); + ei_block_q8_0 *xq = ei_xmalloc((size_t)n * k / 32 * sizeof *xq); + float *row = ei_xmalloc((size_t)k * sizeof(float)), *y = ei_xmalloc((size_t)n * m * sizeof(float)); + for (int r = 0; r < m; r++) { + for (int z = 0; z < k; z++) + row[z] = frand() * (1.0f + (z % 97 == 0) * 6.0f); /* a few outliers per row */ + ei_quantize_row_q8_0(row, wq + (size_t)r * k / 32, k); + } + for (int t = 0; t < n; t++) { + for (int z = 0; z < k; z++) + row[z] = frand() * 3.0f; + ei_quantize_row_q8_0(row, xq + (size_t)t * k / 32, k); + } + /* Repack caches are keyed by tensor address: one live tensor object per case. */ + ei_tensor *wt = ei_xmalloc(sizeof *wt); + *wt = (ei_tensor){.n_dims = 2, .ne = {(uint64_t)k, (uint64_t)m, 1, 1}, .type = EI_T_Q8_0, .data = wq}; + ei_tensor w = *wt; + double times[5]; + for (int rep = 0; rep < 6; rep++) { + double t0 = now_ms(); + ei_gemma2_fast_gemm_q8(ctx, pool, wt, xq, n, y); + if (rep) + times[rep - 1] = now_ms() - t0; /* rep 0 includes the one-time repack */ + } + qsort(times, 5, sizeof *times, cmp_double); + double max_rel = 0; + int checked = 0; + for (int t = 0; t < n; t += n > 64 ? 37 : 1) + for (int r = 0; r < m; r += m > 256 ? 13 : 1) { + double ref = 0, mag = 0; + for (int p = 0; p < k / 32; p++) { + const ei_block_q8_0 *a = wq + (size_t)r * k / 32 + p, *b = xq + (size_t)t * k / 32 + p; + int dot = 0, absdot = 0; + for (int z = 0; z < 32; z++) + dot += a->qs[z] * b->qs[z], absdot += abs(a->qs[z] * b->qs[z]); + double s = (double)ei_fp16_to_fp32(a->d) * ei_fp16_to_fp32(b->d); + ref += dot * s; + mag += absdot * s; + } + double rel = fabs(y[(size_t)t * m + r] - ref) / (mag > 0 ? mag : 1); + if (rel > max_rel) + max_rel = rel; + checked++; + } + double gmacs = (double)n * m * k / (times[2] * 1e6); + printf("%-16s median %8.3f ms %8.1f GMAC/s max_rel_err %.2e (%d outputs checked)\n", cases[c], times[2], + gmacs, max_rel, checked); + if (!(max_rel < 1e-5)) + ok = false; + (void)w; + free(xq), free(row), free(y); /* wq/wt stay alive: the repack cache keys on them */ + } + ei_gemma2_fast_free(ctx); + ei_thread_pool_destroy(pool); + return ok ? 0 : 1; +}