diff --git a/lang/runtime/eg_cosine_batch_strategy_ggml.c b/lang/runtime/eg_cosine_batch_strategy_ggml.c index 73ce6f3..1bc68aa 100644 --- a/lang/runtime/eg_cosine_batch_strategy_ggml.c +++ b/lang/runtime/eg_cosine_batch_strategy_ggml.c @@ -69,12 +69,106 @@ * three strategies. Only the O(n*dim*nq) dot-product matmul — the actual * expensive part — is offloaded to the GPU. * - * Precision: ggml_mul_mat on F32 x F32 inputs computes in F32 on the Metal - * backend (verified: no GGML_PREC_F16 default path applies to F32 inputs; - * see ggml_mul_mat_set_prec in ggml.h, which exists specifically to raise - * precision for lower-than-F32 inputs — ours are already F32 throughout). - * The measured delta vs the CPU double-accumulated oracle is reported - * honestly in the PR body (vindex_bench's BRUTE-GGML line), not assumed. + * ── Precision: the ne11<=8 chunking, and why it is not optional ────────── + * + * The claim in the first version of this file — "ggml_mul_mat on F32 x F32 + * inputs computes in F32 on the Metal backend" — is WRONG, and the 0.9933 + * id-recall it shipped with (vs the hand-rolled kernel's 0.9997) was the + * symptom. ggml-metal has two F32xF32 matmul kernels and picks between them + * purely on ne11 (the number of B rows == our query count): + * + * ne11 <= 8 -> kernel_mul_mv_ext_f32_f32_* / kernel_mul_mv_f32_f32_* + * templated — genuine F32 accumulation. + * ne11 > 8 -> kernel_mul_mm_f32_f32, which is templated + * — i.e. BOTH operands are narrowed + * to F16 and accumulated in simdgroup_half8x8 tiles, even + * though the tensors are GGML_TYPE_F32 on both sides. + * + * (Read it yourself, no guessing — the kernel templates are literal strings + * in the shipped plugin: + * strings $(brew --prefix ggml)/libexec/libggml-metal.so \ + * | grep -E 'host_name\("kernel_mul_m[mv]_f32_f32' + * and the runtime pick is visible with GGML_METAL_DEBUG-style logging as + * "compiling pipeline: base = 'kernel_mul_mm_f32_f32'".) + * + * The previous code issued ONE ggml_mul_mat with ne11 = nq (300 in the + * benchmark), landing squarely on the F16 mul_mm path. Measured on this + * machine (M4 Pro), n=13415 x dim=768 x nq=300, against a CPU double- + * accumulated oracle: + * + * ne11=300 (one mul_mat, the old code) : mean |Δdot| = 1.038e-05 + * ne11=8 (chunked, this code) : mean |Δdot| = 3.863e-09 + * + * — a ~2700x reduction in dot-product error, which is exactly the gap that + * showed up as 0.9933-vs-0.9997 recall. + * + * ggml_mul_mat_set_prec(t, GGML_PREC_F32) does NOT fix this. It was tried: + * the error was bit-identical with and without it (1.038e-05 either way), + * because ggml-metal only consults the prec flag on paths that have an F32 + * variant to switch to, and there is no F32-accumulating mul_mm kernel in + * this build to select. The ONLY lever from outside ggml is ne11. + * + * So: instead of one mul_mat with ne11=nq, we emit ceil(nq/8) mul_mats, each + * over an ne11<=8 ggml_view_2d slice of the same query tensor, all into ONE + * graph and ONE ggml_backend_graph_compute. The node matrix is still uploaded + * exactly once and still read by the GPU as one shared operand — the whole + * point of batch_multi is preserved. + * + * The cost is real, and stated rather than buried. Timing the whole + * batch_multi() call (gather + norms + upload + GPU + scatter) on the real + * shape, median of 15 reps after a discarded warm-up, three separate runs: + * + * unchunked (old, F16 mm) : 13.19 / 13.35 / 14.42 ms -> ~0.044 ms/query + * chunked (this code) : 19.92 / 20.08 / 20.23 ms -> ~0.067 ms/query + * hand-rolled Metal : 17.74 / 17.88 / 17.99 ms -> ~0.060 ms/query + * + * So correctness here costs about +6.7ms per 300-query batch (~1.5x on this + * call), and leaves us ~12% behind the hand-rolled kernel instead of ~35% + * ahead of it. That is not free and should not be sold as free. The reason it + * cannot be recovered inside ggml: an fp32 matmul on Metal has to re-stream + * the whole node matrix once per <=8 queries (38 dispatches x ~41MB here), + * where the F16 mul_mm kernel tiles it in threadgroup memory and reads it far + * fewer times. ggml's Metal backend ships no fp32 TILED matmul, so on this + * backend "fast" and "fp32" are genuinely exclusive — the hand-rolled kernel + * escapes the choice only because it is an fp32 kernel written for this one + * shape. Trading precision back for speed is a one-line env change; trading + * the other way was not available before this commit at all. + * + * 8 is not a magic number we invented — it is ggml-metal's own mul_mm + * threshold, measured by sweeping ne11 and watching both the error and which + * pipeline ggml compiles (9 flips to mul_mm and the error jumps back to + * 1.0e-05 in the same step). EL_GGML_MULMAT_CHUNK overrides it: raise it to + * trade this precision back for throughput, or set it >= nq to reproduce the + * old single-mul_mat behaviour exactly. If a future ggml moves the threshold, + * the worst case is that we silently land back on mul_mm — the same accuracy + * we shipped before, never a correctness break. + * + * ── Cold start: what is and is not ours to fix ─────────────────────────── + * + * The ~7.8s first-call cost reported for the first version of this file is + * NOT this file re-initialising per call (init is, and always was, cached + * behind g_init_attempted below). It is Apple's Metal shader cache missing + * on ggml's embedded metallib — ggml-metal ships ~650 kernels in one + * __ggml_metallib section, and the first newLibraryWithData of it on a given + * machine costs seconds ("ggml_metal_library_init: loaded in 7.670 sec") + * while the driver populates ~/…/C/com.apple.metal/. That cache is keyed on + * the library, not on our binary, and is shared across processes: the very + * next run of a DIFFERENT binary linking the same ggml reports + * "loaded in 0.009 sec". So it is a once-per-machine, per-ggml-version cost, + * not a per-process one, and nothing this file does can avoid it — the + * hand-rolled strategy escapes it only because its shader is two small + * kernels instead of six hundred. + * + * The residual warm init IS ours to look at, and the answer there is "there + * was nothing much to win": ggml_backend_load_all_from_path() dlopens every + * plugin in the directory (three CPU micro-arch variants + BLAS + Metal) when + * we only ever use Metal, so we now load the single Metal plugin instead — + * but measured warm that is 44.7-52.4ms against 46.9-58.9ms, i.e. the same + * number inside noise, because libggml-metal.so's own init dominates. Warm + * ggml init lands at 44-53ms, against 36-117ms for the hand-rolled strategy's + * device+pipeline setup. Cold start was never the real defect here; precision + * was. */ #include "eg_cosine_batch_strategy.h" @@ -82,6 +176,7 @@ #include #include +#include #include #include #include @@ -106,17 +201,92 @@ static const char* eg_ggml_backend_dir(void) { return "/opt/homebrew/opt/ggml/libexec"; } +/* "/libggml-metal.so" in a static buffer. Only ever called once, from + * eg_ggml_ensure_init(), before any thread could race it. */ +static const char* eg_ggml_metal_plugin_path(const char* dir) { + static char buf[1024]; + snprintf(buf, sizeof buf, "%s/libggml-metal.so", dir); + return buf; +} + +/* Largest ne11 (query-batch rows per ggml_mul_mat) that keeps ggml-metal on + * its F32 mul_mv kernels instead of the F16-accumulating mul_mm kernel — see + * the precision discussion in this file's header. EL_GGML_MULMAT_CHUNK + * overrides; a value <= 0 means "use the default". */ +#define EG_GGML_MULMAT_CHUNK_DEFAULT 8 + +static int32_t eg_ggml_mulmat_chunk(void) { + static bool resolved = false; + static int32_t chunk = EG_GGML_MULMAT_CHUNK_DEFAULT; + if (!resolved) { + resolved = true; + const char* s = getenv("EL_GGML_MULMAT_CHUNK"); + if (s && *s) { + long v = strtol(s, NULL, 10); + if (v > 0 && v <= INT32_MAX) chunk = (int32_t)v; + } + } + return chunk; +} + +/* Which ggml device this strategy computes on. GPU (Metal) is the default + * because offloading is the architectural point — the engram's own graph + * traversal and activation spreading are CPU work, and a "GPU" strategy that + * quietly saturates the CPU steals from them. + * + * ACCEL (ggml's BLAS/Accelerate plugin) is reachable here mainly as a + * portability fallback and a diagnostic, and it is documented as MEASURED AND + * REJECTED rather than as a recommendation. In an isolated probe that timed + * only ggml_backend_graph_compute, BLAS looked excellent — 3.4-4.0ms for the + * 300-query batch at mean |Δdot| 1.5e-08, i.e. as fast as the old F16 path and + * far more accurate. End to end on the real store through vindex_bench it does + * not hold up: 0.191 ms/query at id-recall 0.9973, against 0.125-0.142 ms/query + * at 0.9987 for the Metal default. It is dominated on BOTH axes, because the + * isolated probe was not competing with the rest of the batch for the same CPU + * cores and the real call path is. Kept because a machine with no usable Metal + * device still wants a working ggml strategy — not because it is faster. */ +static enum ggml_backend_dev_type eg_ggml_device_type(void) { + const char* s = getenv("EL_GGML_DEVICE"); + if (s && *s) { + if (strcmp(s, "accel") == 0) return GGML_BACKEND_DEVICE_TYPE_ACCEL; + if (strcmp(s, "cpu") == 0) return GGML_BACKEND_DEVICE_TYPE_CPU; + } + return GGML_BACKEND_DEVICE_TYPE_GPU; +} + static bool eg_ggml_ensure_init(void) { if (g_init_attempted) return g_init_ok; g_init_attempted = true; const char* dir = eg_ggml_backend_dir(); - /* dlopen every backend plugin .so in `dir` and register its device(s). - * Never throws; a missing/empty directory just means no devices get - * registered and the lookup below fails cleanly. */ - ggml_backend_load_all_from_path(dir); + const enum ggml_backend_dev_type want = eg_ggml_device_type(); - ggml_backend_dev_t dev = ggml_backend_dev_by_type(GGML_BACKEND_DEVICE_TYPE_GPU); + /* Metal is the only backend this strategy uses by default, so load just + * that one plugin rather than dlopening the whole directory (three CPU + * micro-arch variants + BLAS + Metal here). + * + * Be honest about what this buys: almost nothing in wall time. Measured + * warm, three runs each — load-everything 58.9/46.9/55.1ms, Metal-only + * 52.4/51.2/44.7ms. The cost is dominated by dlopening and initialising + * libggml-metal.so itself, not by the four plugins we skip, so the two + * overlap inside noise. It is kept because registering four device types + * we will never dispatch to is untidy and makes ggml_backend_dev_by_type + * ambiguous, not because it is a speedup — do not cite it as one. + * + * ggml_backend_load() returns NULL for a missing or unloadable path, + * which simply falls through to the broader searches below; it is never + * fatal. Any non-default device needs the full directory scan to find + * its plugin, so skip the fast path there. */ + if (want == GGML_BACKEND_DEVICE_TYPE_GPU) + ggml_backend_load(eg_ggml_metal_plugin_path(dir)); + + ggml_backend_dev_t dev = ggml_backend_dev_by_type(want); + if (!dev) { + /* Non-standard layout, a ggml built with a differently-named Metal + * plugin, or a non-default device: dlopen every plugin in `dir`. */ + ggml_backend_load_all_from_path(dir); + dev = ggml_backend_dev_by_type(want); + } if (!dev) { /* Fall back to ggml's own default search heuristics only if the * explicit path above found nothing — avoids double-registering the @@ -125,7 +295,7 @@ static bool eg_ggml_ensure_init(void) { * symlink vs. its own Cellar-relative guess) in the common case * where the explicit path already worked. */ ggml_backend_load_all(); - dev = ggml_backend_dev_by_type(GGML_BACKEND_DEVICE_TYPE_GPU); + dev = ggml_backend_dev_by_type(want); } if (!dev) return false; @@ -220,42 +390,81 @@ static bool eg_ggml_run(const float* queries, int32_t qdim, int32_t nq, return true; } - /* Pass 2 (GPU via ggml): dot = mul_mat(node_matrix[dim,n_valid], - * queries[dim,nq]) -> dot[n_valid, nq], dot[j*n_valid+i] = dot(node_i,query_j). */ + /* Pass 2 (GPU via ggml): dot[j*n_valid + i] = dot(node_i, query_j), + * computed as ceil(nq/chunk) separate ggml_mul_mat ops over ne11<=chunk + * ggml_view_2d slices of ONE query tensor, all expanded into ONE graph + * and run by ONE ggml_backend_graph_compute. Chunking is what keeps + * ggml-metal on its F32 mul_mv kernels rather than the F16-accumulating + * mul_mm kernel (see this file's header); sharing one graph and one + * t_nodes tensor is what keeps the node matrix uploaded exactly once, + * which is the entire reason batch_multi exists. */ + const int32_t chunk = eg_ggml_mulmat_chunk(); + const int32_t ngroups = (nq + chunk - 1) / chunk; + + /* Tensors held by the context: t_nodes, t_query, plus one view and one + * mul_mat result per group. The graph holds at most one node per view and + * one per mul_mat. Slack on both so a ggml that bookkeeps slightly + * differently cannot silently overflow the arena. */ + const size_t n_tensors = (size_t)2 * (size_t)ngroups + 8; + const size_t graph_size = (size_t)2 * (size_t)ngroups + 16; struct ggml_init_params gp = { - .mem_size = ggml_tensor_overhead() * 8 + ggml_graph_overhead(), + .mem_size = ggml_tensor_overhead() * n_tensors + + ggml_graph_overhead_custom(graph_size, false), .mem_buffer = NULL, .no_alloc = true, }; struct ggml_context* ctx = ggml_init(gp); if (!ctx) { free(valid_orig); free(node_norm_sq); free(node_matrix); free(q_norm_sq); return false; } + struct ggml_tensor** t_dots = (struct ggml_tensor**)malloc((size_t)ngroups * sizeof(*t_dots)); + if (!t_dots) { ggml_free(ctx); free(valid_orig); free(node_norm_sq); free(node_matrix); free(q_norm_sq); return false; } + struct ggml_tensor* t_nodes = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, qdim, n_valid); struct ggml_tensor* t_query = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, qdim, nq); - if (!t_nodes || !t_query) { ggml_free(ctx); free(valid_orig); free(node_norm_sq); free(node_matrix); free(q_norm_sq); return false; } - struct ggml_tensor* t_dot = ggml_mul_mat(ctx, t_nodes, t_query); - if (!t_dot) { ggml_free(ctx); free(valid_orig); free(node_norm_sq); free(node_matrix); free(q_norm_sq); return false; } + struct ggml_cgraph* gf = t_nodes && t_query + ? ggml_new_graph_custom(ctx, graph_size, false) : NULL; + if (!gf) { free(t_dots); ggml_free(ctx); free(valid_orig); free(node_norm_sq); free(node_matrix); free(q_norm_sq); return false; } + + bool built = true; + for (int32_t g = 0; g < ngroups; g++) { + const int32_t start = g * chunk; + const int32_t count = (start + chunk <= nq) ? chunk : (nq - start); + struct ggml_tensor* t_qv = ggml_view_2d(ctx, t_query, qdim, count, + t_query->nb[1], + (size_t)start * t_query->nb[1]); + t_dots[g] = t_qv ? ggml_mul_mat(ctx, t_nodes, t_qv) : NULL; + if (!t_dots[g]) { built = false; break; } + ggml_build_forward_expand(gf, t_dots[g]); + } + if (!built) { free(t_dots); ggml_free(ctx); free(valid_orig); free(node_norm_sq); free(node_matrix); free(q_norm_sq); return false; } struct ggml_backend_buffer* buf = ggml_backend_alloc_ctx_tensors(ctx, g_backend); - if (!buf) { ggml_free(ctx); free(valid_orig); free(node_norm_sq); free(node_matrix); free(q_norm_sq); return false; } + if (!buf) { free(t_dots); ggml_free(ctx); free(valid_orig); free(node_norm_sq); free(node_matrix); free(q_norm_sq); return false; } ggml_backend_tensor_set(t_nodes, node_matrix, 0, (size_t)n_valid * qdim * sizeof(float)); ggml_backend_tensor_set(t_query, queries, 0, (size_t)nq * qdim * sizeof(float)); free(node_matrix); /* uploaded; the packed CPU copy is no longer needed */ - struct ggml_cgraph* gf = ggml_new_graph(ctx); - if (!gf) { ggml_backend_buffer_free(buf); ggml_free(ctx); free(valid_orig); free(node_norm_sq); free(q_norm_sq); return false; } - ggml_build_forward_expand(gf, t_dot); enum ggml_status st = ggml_backend_graph_compute(g_backend, gf); if (st != GGML_STATUS_SUCCESS) { - ggml_backend_buffer_free(buf); ggml_free(ctx); + free(t_dots); ggml_backend_buffer_free(buf); ggml_free(ctx); free(valid_orig); free(node_norm_sq); free(q_norm_sq); return false; } - float* dot = (float*)malloc(ggml_nbytes(t_dot)); - if (!dot) { ggml_backend_buffer_free(buf); ggml_free(ctx); free(valid_orig); free(node_norm_sq); free(q_norm_sq); return false; } - ggml_backend_tensor_get(t_dot, dot, 0, ggml_nbytes(t_dot)); + /* Each group's result is [n_valid, count] contiguous, so reading group g + * into dot + start*n_valid reconstructs exactly the same flat + * dot[j*n_valid + w] layout a single ne11=nq mul_mat would have produced — + * Pass 3 below is unchanged by the chunking. */ + float* dot = (float*)malloc((size_t)n_valid * (size_t)nq * sizeof(float)); + if (!dot) { free(t_dots); ggml_backend_buffer_free(buf); ggml_free(ctx); free(valid_orig); free(node_norm_sq); free(q_norm_sq); return false; } + for (int32_t g = 0; g < ngroups; g++) { + const int32_t start = g * chunk; + const int32_t count = (start + chunk <= nq) ? chunk : (nq - start); + ggml_backend_tensor_get(t_dots[g], dot + (size_t)start * n_valid, 0, + (size_t)count * (size_t)n_valid * sizeof(float)); + } + free(t_dots); /* Pass 3 (CPU): combine dot/(||a||*||b||) per (query,node) pair, scatter * into out_scores at ORIGINAL node indices; every excluded row gets