engram: make the ggml batch-cosine strategy actually compute in fp32
#116 shipped the ggml strategy at 0.9933 id-recall against the CPU oracle while the hand-rolled Metal kernel it replaced scored 0.9997 — a ~150x worse error margin. That was not an inherent property of ggml. It was a usage bug in this file, and this commit fixes it. ggml-metal has two F32xF32 matmul kernels and picks between them purely on ne11, the number of B rows, which for us is the query-batch size: ne11 <= 8 -> kernel_mul_mv_ext_f32_f32_* / kernel_mul_mv_f32_f32_*, templated <float, float> — genuine F32. ne11 > 8 -> kernel_mul_mm_f32_f32, templated <half, half4x4, simdgroup_half8x8, half, half2x4, ...> — BOTH operands narrowed to F16, despite F32 tensors on both sides. The old code issued one ggml_mul_mat with ne11 = nq (300 in the benchmark), landing squarely on the F16 path. The file's own header comment asserted the opposite ("computes in F32 on the Metal backend"); that claim was wrong and is replaced with the measurement. Fix: emit ceil(nq/8) mul_mats over ne11<=8 ggml_view_2d slices of one query tensor, all expanded into ONE graph and one ggml_backend_graph_compute, so the node matrix is still uploaded and shared exactly once. EL_GGML_MULMAT_CHUNK overrides the 8; setting it >= nq reproduces the old behaviour exactly, which is also how the before/after below was measured in a single binary. Measured, real store snapshot, 13415 live embedded nodes, dim=768, 300 real queries, vs the CPU double-accumulated oracle (vindex_bench, offline copy of the store — no live service touched): id-recall same-rank |Δdist| max mean old (ne11=300) 0.9933 6.80e-05 1.43e-05 new (ne11<=8) 0.9987 4.77e-07 9.30e-08 hand-rolled 0.9997 3.58e-07 7.55e-08 ~145x better max error, ~154x better mean — now the same order of magnitude as the hand-rolled kernel rather than 150x off it. The cost is real and is documented rather than buried. Median of 15 reps of the whole batch_multi() call, three runs: 13.2-14.4ms unchunked, 19.9-20.2ms chunked, 17.7-18.0ms hand-rolled. Correctness costs ~+6.7ms per 300-query batch and leaves ggml ~12% behind the hand-rolled kernel instead of ~35% ahead. It cannot be recovered inside ggml: an fp32 matmul on Metal must re-stream the node matrix once per <=8 queries, and ggml's Metal backend ships no fp32 TILED matmul, so "fast" and "fp32" are genuinely exclusive there. Two things that did NOT work, recorded so nobody retries them: - ggml_mul_mat_set_prec(t, GGML_PREC_F32) does nothing here. Error was bit-identical with and without it (1.038e-05 either way) — ggml-metal has no F32-accumulating mul_mm kernel to switch to. ne11 is the only lever. - The ACCEL/BLAS device looked excellent in an isolated compute-only probe (3.4-4.0ms, mean |Δdot| 1.5e-08) but is dominated on BOTH axes end-to-end (0.191 ms/query at 0.9973 recall vs 0.125-0.142 at 0.9987), because the probe was not competing for the same CPU cores the real call path is. It stays reachable via EL_GGML_DEVICE as a no-Metal fallback, labelled as measured-and-rejected, not as a recommendation. Also corrected: the ~7.8s "cold start" blamed on this file is not this file re-initialising per call — init was already cached. It is Apple's shader cache missing on ggml's embedded metallib (~650 kernels), keyed on the library and shared across processes: the first load on a machine reports "loaded in 7.670 sec", the next run of a *different* binary reports 0.009 sec. Once per machine per ggml version, not once per process, and not ours to fix. Warm ggml init is 44-53ms vs 36-117ms for the hand-rolled strategy. Loading only libggml-metal.so instead of every plugin in the directory is kept for tidiness, and explicitly documented as NOT a speedup: 44.7-52.4ms against 46.9-58.9ms, the same number inside noise. The -2.0 sentinel contract is unchanged and re-verified at batch sizes that straddle the chunk boundary (1,7,8,9,16,17,33), plus NULL rows, dim mismatches, zero-norm rows, and an all-invalid population. Notably the old ne11=300 path fails that same check at a 2e-6 cosine tolerance with 2299 mismatches, which is an independent confirmation of the defect.
This commit is contained in:
@@ -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 <float, float> — genuine F32 accumulation.
|
||||
* ne11 > 8 -> kernel_mul_mm_f32_f32, which is templated
|
||||
* <half, half4x4, simdgroup_half8x8, half, half2x4,
|
||||
* simdgroup_half8x8, ...> — 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 <ggml-backend.h>
|
||||
#include <ggml-alloc.h>
|
||||
|
||||
#include <stdint.h>
|
||||
#include <stdio.h>
|
||||
#include <stdlib.h>
|
||||
#include <string.h>
|
||||
@@ -106,17 +201,92 @@ static const char* eg_ggml_backend_dir(void) {
|
||||
return "/opt/homebrew/opt/ggml/libexec";
|
||||
}
|
||||
|
||||
/* "<dir>/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
|
||||
|
||||
Reference in New Issue
Block a user