engram: real Metal batch-cosine kernel, wired into vindex_bench's brute-force oracle
The original brief targeted engram_activate's O(N*D) cosq prescan, but #109 (this branch) already retires that loop algorithmically (HNSW seed selection + lazy memoized cosine) — GPU-accelerating a loop being deleted isn't real work, so that target was dropped rather than forced. Re-investigated for a genuine remaining GPU-shaped call site (not a manufactured one): HNSW insert's candidate-list distance work is bounded- degree (M=24-48) and sequential/adaptive — too fine-grained for a GPU dispatch to pay off. No O(N^2) pairwise cosine pass exists (dedup only checks the K=8 already-selected seed slots). No concurrent multi-query traffic exists (server.el: the soul's curiosity loop is a single in-process caller). vindex_bench.c's brute_topk — the correctness oracle this same PR adds to validate HNSW recall — is the one real, unforced fit: genuine 1-query-vs-N-vectors, embarrassingly parallel, no adaptivity. Adds: - eg_cosine_batch.metal: batched cosine kernel, single- and multi-query variants, same -2.0 dim-mismatch/zero-norm sentinel as eg_cosine(). - eg_metal_cosine.h/.m: C-callable Objective-C bridge. Lazy one-time device/pipeline init, MTLResourceStorageModeShared buffers, returns false on ANY failure so callers fall back to the scalar CPU loop unconditionally — never partial, never throws. - eg_metal_cosine_stub.c: zero-dependency CPU-only implementation for non-Darwin builds (Linux CI) — same symbols, always returns false, no #ifdef needed at any call site. - build_vindex_bench.sh: one-command build, real bridge + Metal frameworks on Darwin, stub everywhere else. vindex_bench.c: brute_topk_metal / brute_topk_metal_batch call the bridge, falling back to the existing CPU brute_topk on any failure or EL_METAL_COSINE=0. The multi-query batched path exists because the first version (one GPU call per query) measured SLOWER than CPU at N~13.7k — it re-uploaded the full N*D matrix every query. Fixed by uploading the matrix once per query batch. Measured against a real nsbx-sandboxed clone of the live store (never :8742/:7770), 13,671 real embedded nodes, dim=768, 300 real queries: BRUTE-FORCE (CPU): 2.013 ms/query BRUTE-METAL (GPU): 0.117 ms/query (17.2x) id-recall vs CPU oracle: 0.9990 over 300 queries same-rank |Δdist|: max 2.98e-07, mean 7.53e-08 (float32 rounding, not a bug) Synthetic scaling sweep (13k -> 50k nodes, same dim/queries) shows the GPU speedup holding (~11x) as N grows toward the mathematical-foundations doc's 1.3M-node target, with CPU brute-force cost growing linearly as expected. Not wired into engram_activate or the daemon build (nsbx's _build_binary) — vindex_bench is a standalone offline tool, not part of the request-serving binary, so no engram_activate/server-latency claim is made here. The bridge is a reusable primitive (single eg_cosine_batch_metal + batched eg_cosine_batch_metal_multi) other call sites can adopt later without re-deriving any of this. Based on feat/reframe-region-setop (PR #109), not dev directly: the only genuine batch-cosine call site (vindex_bench.c) exists solely on this branch. Flagged explicitly in the PR description as a deliberate deviation from the original "base off dev" instruction.
This commit is contained in:
bigmerge
@@ -0,0 +1,156 @@
|
||||
/* eg_cosine_batch.metal — batched cosine similarity, one query vs N node vectors.
|
||||
*
|
||||
* GPU-shaped counterpart to eg_cosine() in el_runtime.c: same math, same
|
||||
* dim-mismatch/zero-norm sentinel (-2.0), applied to N independent rows in
|
||||
* parallel instead of one pair at a time in a CPU loop.
|
||||
*
|
||||
* Semantics MUST match eg_cosine() exactly:
|
||||
* - inputs are raw, UNNORMALIZED vectors (nomic-embed-text magnitudes are
|
||||
* not 1.0) — this kernel computes the full dot/(|a|*|b|) cosine, not a
|
||||
* plain dot product.
|
||||
* - a node whose declared dim differs from the query dim, or whose norm is
|
||||
* zero, scores exactly -2.0 (below any valid cosine in [-1,1]), so a
|
||||
* caller doing `if (score > threshold)` behaves identically whether the
|
||||
* scalar or the batched path filled the array.
|
||||
*
|
||||
* Precision: Apple GPUs do not support double in Metal Shading Language —
|
||||
* everything here is float32. eg_cosine accumulates in CPU double, but its
|
||||
* *inputs* are float32 embeddings, so the achievable precision ceiling is
|
||||
* bounded by the input data regardless of accumulator width. To keep the
|
||||
* float32 reduction from drifting relative to the double-accumulated CPU
|
||||
* result across dim=768 terms, each thread accumulates with 4 independent
|
||||
* partial sums (unrolled) rather than one running scalar — the same
|
||||
* error-reduction trick already used by the CPU brute-force loop in
|
||||
* vindex_bench.c. The measured float-vs-double delta is reported in the PR
|
||||
* description; this is not assumed to be "close enough" without measurement.
|
||||
*/
|
||||
#include <metal_stdlib>
|
||||
using namespace metal;
|
||||
|
||||
/* Per-dispatch invariants. `dim` is the query's dimensionality — the
|
||||
* dimensionality every comparable node vector must match. */
|
||||
struct EgCosineParams {
|
||||
uint n; /* number of node rows */
|
||||
uint dim; /* vector width (both query and node rows are `dim` wide in
|
||||
* the packed buffer; node_dims[] carries each node's REAL
|
||||
* embedded dim for the mismatch check) */
|
||||
};
|
||||
|
||||
/* One thread per node row. node_matrix is n*dim floats, row-major, packed at
|
||||
* `dim` stride regardless of a row's real dim (the CPU side zero-pads or
|
||||
* skips packing rows that don't match — see eg_cosine_batch_metal in
|
||||
* eg_metal_cosine.m for the exact packing contract). node_dims[i] is the
|
||||
* node's true emb_dim, used only for the mismatch sentinel — never used to
|
||||
* index, since every row is packed at uniform `dim` stride. */
|
||||
kernel void eg_cosine_batch_kernel(
|
||||
device const float* query [[buffer(0)]],
|
||||
device const float* node_matrix [[buffer(1)]],
|
||||
device const int* node_dims [[buffer(2)]],
|
||||
constant EgCosineParams& p [[buffer(3)]],
|
||||
device float* out_scores [[buffer(4)]],
|
||||
uint gid [[thread_position_in_grid]])
|
||||
{
|
||||
if (gid >= p.n) return;
|
||||
|
||||
if (node_dims[gid] != int(p.dim)) {
|
||||
out_scores[gid] = -2.0f;
|
||||
return;
|
||||
}
|
||||
|
||||
device const float* row = node_matrix + (uint64_t)gid * (uint64_t)p.dim;
|
||||
|
||||
/* 4-way partial accumulation — same shape as vindex_bench.c's brute_topk
|
||||
* unroll, done here for float32 accuracy rather than raw throughput. */
|
||||
float dot0 = 0.0f, dot1 = 0.0f, dot2 = 0.0f, dot3 = 0.0f;
|
||||
float na0 = 0.0f, na1 = 0.0f, na2 = 0.0f, na3 = 0.0f;
|
||||
float nb0 = 0.0f, nb1 = 0.0f, nb2 = 0.0f, nb3 = 0.0f;
|
||||
|
||||
uint d = 0;
|
||||
uint dim4 = p.dim & ~3u;
|
||||
for (; d < dim4; d += 4) {
|
||||
float a0 = row[d], b0 = query[d];
|
||||
float a1 = row[d+1], b1 = query[d+1];
|
||||
float a2 = row[d+2], b2 = query[d+2];
|
||||
float a3 = row[d+3], b3 = query[d+3];
|
||||
dot0 += a0*b0; dot1 += a1*b1; dot2 += a2*b2; dot3 += a3*b3;
|
||||
na0 += a0*a0; na1 += a1*a1; na2 += a2*a2; na3 += a3*a3;
|
||||
nb0 += b0*b0; nb1 += b1*b1; nb2 += b2*b2; nb3 += b3*b3;
|
||||
}
|
||||
float dot = (dot0 + dot1) + (dot2 + dot3);
|
||||
float na = (na0 + na1) + (na2 + na3);
|
||||
float nb = (nb0 + nb1) + (nb2 + nb3);
|
||||
for (; d < p.dim; d++) {
|
||||
float a = row[d], b = query[d];
|
||||
dot += a*b; na += a*a; nb += b*b;
|
||||
}
|
||||
|
||||
if (na <= 0.0f || nb <= 0.0f) {
|
||||
out_scores[gid] = -2.0f;
|
||||
return;
|
||||
}
|
||||
out_scores[gid] = dot / sqrt(na * nb);
|
||||
}
|
||||
|
||||
/* ── multi-query variant ──────────────────────────────────────────────────
|
||||
* Same per-pair math as eg_cosine_batch_kernel, but amortizes ONE upload of
|
||||
* node_matrix (the expensive part at real store size — 13k*768 floats is
|
||||
* ~42MB) across `nq` queries instead of re-uploading it once per query.
|
||||
* Measured need: a naive one-query-at-a-time loop calling the single-query
|
||||
* kernel nq times was SLOWER than the CPU oracle at N≈13.7k (re-gather +
|
||||
* re-upload dominated the actual compute) — this is the fix, not a
|
||||
* hypothetical optimization.
|
||||
*
|
||||
* 2D grid: x = node index [0,n), y = query index [0,nq). out_scores is
|
||||
* nq*n, row-major by query (out_scores[qid*n + nid]). */
|
||||
struct EgCosineMultiParams { uint n; uint dim; uint nq; };
|
||||
|
||||
kernel void eg_cosine_batch_multi_kernel(
|
||||
device const float* queries [[buffer(0)]], /* nq*dim */
|
||||
device const float* node_matrix [[buffer(1)]], /* n*dim */
|
||||
device const int* node_dims [[buffer(2)]], /* n */
|
||||
constant EgCosineMultiParams& p [[buffer(3)]],
|
||||
device float* out_scores [[buffer(4)]], /* nq*n */
|
||||
uint2 gid [[thread_position_in_grid]])
|
||||
{
|
||||
uint nid = gid.x, qid = gid.y;
|
||||
if (nid >= p.n || qid >= p.nq) return;
|
||||
|
||||
uint64_t out_idx = (uint64_t)qid * (uint64_t)p.n + (uint64_t)nid;
|
||||
|
||||
if (node_dims[nid] != int(p.dim)) {
|
||||
out_scores[out_idx] = -2.0f;
|
||||
return;
|
||||
}
|
||||
|
||||
device const float* row = node_matrix + (uint64_t)nid * (uint64_t)p.dim;
|
||||
device const float* query = queries + (uint64_t)qid * (uint64_t)p.dim;
|
||||
|
||||
float dot0 = 0.0f, dot1 = 0.0f, dot2 = 0.0f, dot3 = 0.0f;
|
||||
float na0 = 0.0f, na1 = 0.0f, na2 = 0.0f, na3 = 0.0f;
|
||||
float nb0 = 0.0f, nb1 = 0.0f, nb2 = 0.0f, nb3 = 0.0f;
|
||||
|
||||
uint d = 0;
|
||||
uint dim4 = p.dim & ~3u;
|
||||
for (; d < dim4; d += 4) {
|
||||
float a0 = row[d], b0 = query[d];
|
||||
float a1 = row[d+1], b1 = query[d+1];
|
||||
float a2 = row[d+2], b2 = query[d+2];
|
||||
float a3 = row[d+3], b3 = query[d+3];
|
||||
dot0 += a0*b0; dot1 += a1*b1; dot2 += a2*b2; dot3 += a3*b3;
|
||||
na0 += a0*a0; na1 += a1*a1; na2 += a2*a2; na3 += a3*a3;
|
||||
nb0 += b0*b0; nb1 += b1*b1; nb2 += b2*b2; nb3 += b3*b3;
|
||||
}
|
||||
float dot = (dot0 + dot1) + (dot2 + dot3);
|
||||
float na = (na0 + na1) + (na2 + na3);
|
||||
float nb = (nb0 + nb1) + (nb2 + nb3);
|
||||
for (; d < p.dim; d++) {
|
||||
float a = row[d], b = query[d];
|
||||
dot += a*b; na += a*a; nb += b*b;
|
||||
}
|
||||
|
||||
if (na <= 0.0f || nb <= 0.0f) {
|
||||
out_scores[out_idx] = -2.0f;
|
||||
return;
|
||||
}
|
||||
out_scores[out_idx] = dot / sqrt(na * nb);
|
||||
}
|
||||
Reference in New Issue
Block a user