engram: make the ggml batch-cosine strategy actually compute in fp32 (recall 0.9933 -> 0.9987) #121
Reference in New Issue
Block a user
Delete Branch "improve/ggml-cosine-fp32-and-init"
Deleting a branch is permanent. Although the deleted branch may continue to exist for a short time before it actually gets removed, it CANNOT be undone in most cases. Continue?
Follow-up to #116. Base is
feat/engram-ggml-cosine-batch(#116's branch), notdev— #116 is still open, sodevdoes not containeg_cosine_batch_strategy_ggml.cyet. Targeting #116's branch keeps this diff to just the improvement; retarget todevif #116 merges first. Nothing in #116's history is rebased, amended or rewritten.Why
Will, on #116's own numbers (ggml recall 0.9933 vs the hand-rolled kernel's 0.9997): "we can keep ggml, but why do it worse?"
Answer: we don't have to. The precision gap was a real, fixable usage bug in how this file called ggml — not a ggml limitation. The cold-start gap was not a bug at all. Both are settled below with measurements, including the fixes that didn't work.
The precision bug
ggml-metalhas two F32xF32 matmul kernels and picks between them purely onne11— the number of B rows, which for us is the query-batch size:ne11<= 8kernel_mul_mv_ext_f32_f32_*<float, float>— genuine F32> 8kernel_mul_mm_f32_f32<half, half4x4, simdgroup_half8x8, half, half2x4, ...>— both operands narrowed to F16Both operands are narrowed to F16 even though the tensors are
GGML_TYPE_F32on both sides. Verify it yourself, no guessing — the kernel templates are literal strings in the shipped plugin:#116issued oneggml_mul_matwithne11 = nq(300 in the benchmark), landing squarely on the F16 path. ggml confirms the pick at runtime:The file's own header comment asserted the opposite ("computes in F32 on the Metal backend"). That claim was wrong; it is replaced with the measurement.
Fix
Emit
ceil(nq/8)mul_mats overne11<=8ggml_view_2dslices of one query tensor, all expanded into one graph and oneggml_backend_graph_compute— so the node matrix is still uploaded and shared exactly once, which is the entire reasonbatch_multiexists.EL_GGML_MULMAT_CHUNKoverrides the 8; setting it>= nqreproduces the old behaviour exactly, which is how the before/after below was measured in a single binary.8is not invented — it is ggml-metal's own threshold, found by sweepingne11and watching both the error and which pipeline ggml compiles.9flips tomul_mmand the error jumps back in the same step:Results — real store, offline snapshot
vindex_bench store <snapshot> 768 300 10 64against an offline copy of the store (13,415 live embedded nodes, dim=768, 300 real queries, CPU double-accumulated oracle). No live service on:8742/:7770was touched.ne11=300)ne11<=8)~145x better max error, ~154x better mean — now the same order of magnitude as the hand-rolled kernel instead of 150x off it. These three lines are bit-identical on every repeat run.
The cost, stated rather than buried
The
ms/querycolumn invindex_benchis too noisy on this machine to carry a claim (the untouched hand-rolled strategy swings 0.064–0.131 across runs). So speed was measured separately: median of 15 reps of the wholebatch_multi()call — gather + norms + upload + GPU + scatter — after a discarded warm-up, three separate runs:Correctness costs ~+6.7ms per 300-query batch (~1.5x on this call), and leaves ggml ~12% behind the hand-rolled kernel instead of ~35% ahead of it.
This is not recoverable inside ggml. An fp32 matmul on Metal must re-stream the whole node matrix once per
<=8queries (38 dispatches x ~41MB here), where the F16mul_mmkernel 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.Decision for you, not taken here: the factory still prefers
ggmlfirst, per your original directive. On this hardware and shape ggml is now measurably behind the hand-rolled kernel on both axes (0.067 vs 0.060 ms/query, 0.9987 vs 0.9997). Flipping the default preference order is a policy call, so it is left alone and flagged rather than changed unilaterally.Two things that did NOT work
Recorded so nobody retries them.
ggml_mul_mat_set_prec(t, GGML_PREC_F32)does nothing here. Tried it; error was bit-identical with and without —1.038e-05either way. ggml-metal has no F32-accumulatingmul_mmkernel to switch to, so the flag has nowhere to go.ne11is the only lever.|Δdot|1.5e-08, i.e. as fast as the old F16 path and far more accurate. End-to-end on the real store it is dominated on both axes: 0.191 ms/query at 0.9973 recall, vs 0.125–0.142 at 0.9987 for the Metal default. The probe wasn't competing for the same CPU cores the real call path is. It stays reachable viaEL_GGML_DEVICE=accelas a no-Metal fallback, labelled measured-and-rejected.Cold start — not a bug, and not ours
The ~7.8s first-call cost attributed to this file is not this file re-initialising per call. Init was already cached behind
g_init_attempted(and always was — the hypothesis that it re-ranggml_backend_load_all_from_pathevery call is false; see the original lines 109–111).It is Apple's Metal shader cache missing on ggml's embedded metallib (~650 kernels in one
__ggml_metallibsection). That cache is keyed on the library, not our binary, and is shared across processes:Corroborated by the cache itself —
~/…/C/com.apple.metal/32024/libraries3.datagrew to 48MB at the moment of the 7.67s run. So: once per machine per ggml version, not once per process, 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.Warm ggml init is 44–53ms, against 36–117ms for the hand-rolled strategy's device+pipeline setup — comparable.
Loading only
libggml-metal.soinstead of every plugin in the directory is kept for tidiness and explicitly documented as not a speedup: 44.7/51.2/52.4ms against 46.9/55.1/58.9ms, the same number inside noise.Contract
The
-2.0sentinel contract is unchanged and re-verified at batch sizes straddling the chunk boundary (1, 7, 8, 9, 16, 17, 33), plus NULL rows, dim mismatches, zero-norm rows, and an all-invalid population — all pass, under the default and underEL_GGML_MULMAT_CHUNK=1/300andEL_GGML_DEVICE=accel/cpu.Independent confirmation of the defect: the old
ne11=300path fails that same check at a 2e-6 cosine tolerance with 2299 mismatches, while the new default passes with 0.Scope
One file:
lang/runtime/eg_cosine_batch_strategy_ggml.c.eg_cosine_batch.metalandeg_cosine_batch_strategy_metal_hand.mare untouched — read only as reference.