perf(retain): accelerate within-batch semantic link calculation (#3977)

Cuts the within-batch semantic link pass to float32 and one reused buffer. The batch was widened to float64, but PackedEmbedding is array("f") and pgvector's vector column stores float32, so the extra 32 bits were padding nothing downstream could read. Dropping to float32 halves the working set and puts BLAS on SGEMM; normalising in place, deriving validity from the row norms instead of an (n, dim) isfinite mask, and reusing one similarity buffer across blocks remove three further copies. Peak transient falls 74-86% (at 5,000 facts, 235 MB -> 48 MB). argpartition replaces a full sort that existed only to discard all but top_k, and the self-link mask and score unboxing move out of the Python loop: 1.7-2.2x on a realistic clustered batch, up to 4.8x when nearly every pair clears the threshold. Norms are accumulated in float64 via einsum, since a float32 sum of 1536 squares overflows above ~1e19 and flushes to zero below ~1e-22. The batch is copied with np.array rather than aliased with asarray, as it is now normalised in place. Verified against the float64 implementation across 120 randomised batches plus NaN/inf/zero embeddings, degenerate magnitudes and all-ties: identical link pairs, scores within 1e-6.
2026-09-14 19:31:49 +08:00 · 2026-09-08 17:38:19 +08:00
parent 3e1d47fdc6
commit b045794817
5 changed files with 319 additions and 23 deletions
@@ -743,14 +743,29 @@ def compute_semantic_links_within_batch(
    if len(unit_ids) < 2:
        return []

+    n_units = len(unit_ids)
    links = []
-    new_embeddings_matrix = np.asarray(embeddings, dtype=float)
-    norms = np.linalg.norm(new_embeddings_matrix, axis=1)
-    valid_embeddings = np.isfinite(new_embeddings_matrix).all(axis=1) & np.isfinite(norms) & (norms > 0)
-    normalized_embeddings = np.zeros_like(new_embeddings_matrix)
-    normalized_embeddings[valid_embeddings] = (
-        new_embeddings_matrix[valid_embeddings] / norms[valid_embeddings, np.newaxis]
+    # float32, not float64: `PackedEmbedding` is already `array("f")` and pgvector's `vector`
+    # column stores float32, so the doubles the old `dtype=float` produced were padding that
+    # nothing downstream could use -- they only doubled the working set and pushed BLAS off
+    # SGEMM onto DGEMM. `np.array` (not `asarray`) because this buffer is normalised in place
+    # below and must not alias an ndarray the caller still owns.
+    normalized_embeddings = np.array(embeddings, dtype=np.float32)
+    # Accumulate the norms in float64. The vectors are float32, but summing 1536 squares in
+    # float32 overflows to inf above ~1e19 and flushes to zero below ~1e-22, which would drop
+    # those rows as "invalid" when float64 handled them fine. `einsum` keeps the wide
+    # accumulator without materialising an (n, dim) float64 copy of the batch.
+    norms = np.sqrt(np.einsum("ij,ij->i", normalized_embeddings, normalized_embeddings, dtype=np.float64))
+    # A non-finite component poisons its own row norm, so the norm check alone catches NaN and
+    # inf rows -- no need for an (n, dim) `isfinite` mask over the whole batch.
+    valid_embeddings = np.isfinite(norms) & (norms > 0)
+    np.divide(
+        normalized_embeddings,
+        norms[:, np.newaxis].astype(np.float32),
+        out=normalized_embeddings,
+        where=valid_embeddings[:, np.newaxis],
    )
+    normalized_embeddings[~valid_embeddings] = 0.0

    # One matrix product per block of rows, rather than one per unit against a
    # freshly gathered copy of every other unit. `normalized[others]` is advanced
@@ -761,26 +776,48 @@ def compute_semantic_links_within_batch(
    # O(n^2 * dim) dot products either way; this hands them to BLAS in one call and
    # keeps the transient at one block of similarity rows.
    block_rows = _SEMANTIC_WITHIN_BATCH_BLOCK_ROWS
-    for start in range(0, len(unit_ids), block_rows):
-        stop = min(start + block_rows, len(unit_ids))
-        block_similarities = normalized_embeddings[start:stop] @ normalized_embeddings.T
+    # One (block_rows, n) buffer for the whole sweep. At 36K facts each block of
+    # similarities is 37 MB, and allocating and freeing that once per block is churn
+    # the allocator does not need to see.
+    similarity_buffer = np.empty((min(block_rows, n_units), n_units), dtype=np.float32)
+    invalid_columns = np.flatnonzero(~valid_embeddings)
+    for start in range(0, n_units, block_rows):
+        stop = min(start + block_rows, n_units)
+        block_similarities = similarity_buffer[: stop - start]
+        np.matmul(normalized_embeddings[start:stop], normalized_embeddings.T, out=block_similarities)
        # A unit with an unusable embedding is neither a source nor a target.
-        block_similarities[:, ~valid_embeddings] = -np.inf
+        if invalid_columns.size:
+            block_similarities[:, invalid_columns] = -np.inf
+        # Never link a unit to itself: row `i` of the block is unit `start + i`.
+        diagonal = np.arange(stop - start)
+        block_similarities[diagonal, start + diagonal] = -np.inf

-        for i in range(start, stop):
-            if not valid_embeddings[i]:
+        for local_index, unit_index in enumerate(range(start, stop)):
+            if not valid_embeddings[unit_index]:
                continue

-            similarities = block_similarities[i - start]
-            similarities[i] = -np.inf  # never link a unit to itself
-
+            similarities = block_similarities[local_index]
            above_threshold = np.where(similarities >= threshold)[0]
-            if len(above_threshold) > 0:
-                sorted_indices = above_threshold[np.argsort(-similarities[above_threshold])][:top_k]
-                for other_idx in sorted_indices:
-                    other_id = unit_ids[other_idx]
-                    similarity = float(min(1.0, max(0.0, similarities[other_idx])))
-                    links.append((unit_ids[i], other_id, "semantic", similarity, None))
+            candidate_count = len(above_threshold)
+            if candidate_count == 0:
+                continue
+
+            if candidate_count > top_k:
+                # Introselect the top k in O(candidates), then sort only those k, rather
+                # than sorting every candidate to throw all but k of them away.
+                candidate_scores = -similarities[above_threshold]
+                top_partition = np.argpartition(candidate_scores, top_k)[:top_k]
+                neighbours = above_threshold[top_partition[np.argsort(candidate_scores[top_partition])]]
+            elif candidate_count > 1:
+                neighbours = above_threshold[np.argsort(-similarities[above_threshold])]
+            else:
+                neighbours = above_threshold
+
+            from_id = unit_ids[unit_index]
+            # One C-level pass to clamp and unbox, instead of boxing each score on its own.
+            scores = np.clip(similarities[neighbours], 0.0, 1.0).tolist()
+            for other_index, similarity in zip(neighbours, scores):
+                links.append((from_id, unit_ids[other_index], "semantic", similarity, None))

    return links

@@ -210,7 +210,7 @@ class TestWithinBatchSemanticLinks:
        legacy = _legacy_semantic_links_within_batch(unit_ids, embeddings, top_k=8, threshold=0.1)

        assert [(link[0], link[1], link[2]) for link in new] == [(link[0], link[1], link[2]) for link in legacy]
-        assert [link[3] for link in new] == pytest.approx([link[3] for link in legacy], abs=1e-12)
+        assert [link[3] for link in new] == pytest.approx([link[3] for link in legacy], abs=1e-6)

    def test_matches_across_block_boundaries(self, monkeypatch):
        """Blocking must not change which units are compared, only when."""
@@ -223,7 +223,44 @@ class TestWithinBatchSemanticLinks:
        legacy = _legacy_semantic_links_within_batch(unit_ids, embeddings, top_k=5, threshold=0.0)

        assert [(link[0], link[1]) for link in new] == [(link[0], link[1]) for link in legacy]
-        assert [link[3] for link in new] == pytest.approx([link[3] for link in legacy], abs=1e-12)
+        assert [link[3] for link in new] == pytest.approx([link[3] for link in legacy], abs=1e-6)
+
+    def test_matches_the_float64_reference_on_degenerate_magnitudes(self):
+        """Norms are accumulated in float64: a float32 sum of squares would overflow to inf
+        above ~1e19 and flush to zero below ~1e-22, silently dropping those rows."""
+        rng = np.random.default_rng(5)
+        for scale in (1e-30, 1e-23, 1e20, 1e25):
+            embeddings = (rng.normal(size=(6, 8)) * scale).astype(np.float32)
+            unit_ids = [f"u{i}" for i in range(6)]
+
+            new = compute_semantic_links_within_batch(unit_ids, embeddings, top_k=3, threshold=0.0)
+            legacy = _legacy_semantic_links_within_batch(unit_ids, embeddings, top_k=3, threshold=0.0)
+
+            assert [(link[0], link[1]) for link in new] == [(link[0], link[1]) for link in legacy], scale
+            assert [link[3] for link in new] == pytest.approx([link[3] for link in legacy], abs=1e-6)
+
+    def test_does_not_mutate_the_caller_s_embeddings(self):
+        """Normalisation runs in place, so the batch must be copied off the caller's buffer."""
+        rng = np.random.default_rng(6)
+        embeddings = rng.normal(size=(5, 8)).astype(np.float32)
+        before = embeddings.copy()
+
+        compute_semantic_links_within_batch([f"u{i}" for i in range(5)], embeddings, top_k=3, threshold=0.0)
+
+        assert np.array_equal(embeddings, before)
+
+    def test_top_k_is_honoured_when_every_candidate_ties(self, monkeypatch):
+        """Duplicate facts give every pair the same score; introselect must still emit exactly
+        top_k neighbours per source rather than tripping over the ties."""
+        monkeypatch.setattr(link_utils, "_SEMANTIC_WITHIN_BATCH_BLOCK_ROWS", 4)
+        unit_ids = [f"u{i}" for i in range(12)]
+        embeddings = [[0.5, 0.25, 0.125, 0.0625]] * 12
+
+        links = compute_semantic_links_within_batch(unit_ids, embeddings, top_k=5, threshold=0.5)
+
+        assert len(links) == 12 * 5
+        assert all(sum(1 for link in links if link[0] == unit_id) == 5 for unit_id in unit_ids)
+        assert not [link for link in links if link[0] == link[1]]

    def test_invalid_embeddings_stay_excluded_across_blocks(self, monkeypatch):
        monkeypatch.setattr(link_utils, "_SEMANTIC_WITHIN_BATCH_BLOCK_ROWS", 2)
@@ -0,0 +1,206 @@
+"""Microbenchmark for ``compute_semantic_links_within_batch`` (retain phase 2).
+
+Timing and memory are measured in *separate* passes on purpose: ``tracemalloc`` taxes every
+allocation, and it taxes the allocation-heavy baseline hardest, so timing a traced run
+understates the speedup by up to 4x.
+
+Wall time is the best of N runs, and the before/after rounds are *interleaved* rather than run
+back to back. On a thermally loaded laptop the same measurement drifts by more than 2x over a
+few minutes, so running all of "before" and then all of "after" attributes that drift to the
+code. Interleaving cancels it. Treat the speedup column as a ratio with roughly +/-20% of play
+in it; the memory column is exact and repeats to the byte.
+"""
+
+import argparse
+import gc
+import time
+import tracemalloc
+from array import array
+from collections.abc import Callable, Sequence
+from dataclasses import dataclass
+
+import numpy as np
+from hindsight_api.engine.retain.link_utils import compute_semantic_links_within_batch
+from rich.console import Console
+from rich.table import Table
+
+console = Console()
+
+# Frozen snapshot of the implementation before #3977 -- float64 throughout, full argsort per
+# row. Kept verbatim so the table has a real "before" column; it is not imported from the
+# engine on purpose, since the point is to compare against code that no longer exists.
+_BASELINE_BLOCK_ROWS = 256
+
+
+def _baseline(unit_ids: list[str], embeddings: Sequence, top_k: int = 50, *, threshold: float) -> list[tuple]:
+    if len(unit_ids) < 2:
+        return []
+    links = []
+    matrix = np.asarray(embeddings, dtype=float)
+    norms = np.linalg.norm(matrix, axis=1)
+    valid = np.isfinite(matrix).all(axis=1) & np.isfinite(norms) & (norms > 0)
+    normalized = np.zeros_like(matrix)
+    normalized[valid] = matrix[valid] / norms[valid, np.newaxis]
+    for start in range(0, len(unit_ids), _BASELINE_BLOCK_ROWS):
+        stop = min(start + _BASELINE_BLOCK_ROWS, len(unit_ids))
+        block = normalized[start:stop] @ normalized.T
+        block[:, ~valid] = -np.inf
+        for i in range(start, stop):
+            if not valid[i]:
+                continue
+            similarities = block[i - start]
+            similarities[i] = -np.inf
+            above = np.where(similarities >= threshold)[0]
+            if len(above) > 0:
+                for other in above[np.argsort(-similarities[above])][:top_k]:
+                    score = float(min(1.0, max(0.0, similarities[other])))
+                    links.append((unit_ids[i], unit_ids[other], "semantic", score, None))
+    return links
+
+
+@dataclass(frozen=True)
+class Workload:
+    """One synthetic similarity regime to measure the pass against."""
+
+    threshold: float
+    description: str
+
+
+# "clustered" is the one that resembles a real retain batch: topical clusters, so a realistic
+# fraction of pairs clears the default 0.7 threshold. "sparse" deliberately produces no links at
+# all, which isolates the matrix product from the selection and materialisation work.
+WORKLOADS = {
+    "sparse": Workload(0.7, "uncorrelated facts, near-zero links -- isolates the matrix product"),
+    "clustered": Workload(0.7, "topical clusters at the default threshold -- the realistic retain batch"),
+    "dense": Workload(0.3, "everything similar -- stresses candidate extraction and top-k"),
+}
+
+
+@dataclass(frozen=True)
+class Batch:
+    unit_ids: list[str]
+    embeddings: list[array]
+
+
+def _make_batch(n: int, kind: str, dim: int, seed: int = 0) -> Batch:
+    """Embeddings as ``array("f")`` -- the ``PackedEmbedding`` the retain pipeline carries."""
+    rng = np.random.default_rng(seed)
+    if kind == "sparse":
+        matrix = rng.normal(size=(n, dim))
+    elif kind == "clustered":
+        centers = rng.normal(size=(max(2, n // 25), dim))
+        matrix = centers[rng.integers(0, len(centers), n)] + rng.normal(scale=0.55, size=(n, dim))
+    else:
+        matrix = rng.random(size=(n, dim))
+    return Batch(
+        unit_ids=[f"u{i}" for i in range(n)],
+        embeddings=[array("f", row.tolist()) for row in matrix.astype(np.float32)],
+    )
+
+
+@dataclass(frozen=True)
+class Round:
+    elapsed_ms: float
+    links: int
+
+
+@dataclass(frozen=True)
+class Timings:
+    before_ms: float
+    after_ms: float
+    links: int
+
+    @property
+    def speedup(self) -> float:
+        return self.before_ms / self.after_ms
+
+
+def _one_round(fn: Callable, batch: Batch, threshold: float, top_k: int) -> Round:
+    gc.collect()
+    started = time.perf_counter()
+    links = fn(batch.unit_ids, batch.embeddings, top_k=top_k, threshold=threshold)
+    elapsed = (time.perf_counter() - started) * 1000
+    count = len(links)
+    del links
+    return Round(elapsed_ms=elapsed, links=count)
+
+
+def _interleaved_best(
+    before: Callable, after: Callable, batch: Batch, threshold: float, top_k: int, repeats: int
+) -> Timings:
+    """Alternate the two implementations so CPU frequency drift hits both equally."""
+    for fn in (before, after):
+        fn(batch.unit_ids, batch.embeddings, top_k=top_k, threshold=threshold)  # warm BLAS and the allocator
+    best_before = best_after = float("inf")
+    links = 0
+    for _ in range(repeats):
+        best_before = min(best_before, _one_round(before, batch, threshold, top_k).elapsed_ms)
+        round_after = _one_round(after, batch, threshold, top_k)
+        best_after = min(best_after, round_after.elapsed_ms)
+        links = round_after.links
+    return Timings(before_ms=best_before, after_ms=best_after, links=links)
+
+
+def _peak_mb(fn: Callable, batch: Batch, threshold: float, top_k: int) -> float:
+    gc.collect()
+    tracemalloc.start()
+    links = fn(batch.unit_ids, batch.embeddings, top_k=top_k, threshold=threshold)
+    _, peak = tracemalloc.get_traced_memory()
+    tracemalloc.stop()
+    del links
+    return peak / (1024 * 1024)
+
+
+def run_benchmark(
+    n_values: Sequence[int],
+    dim: int,
+    top_k: int,
+    repeats: int,
+    workloads: Sequence[str],
+) -> None:
+    for kind in workloads:
+        workload = WORKLOADS[kind]
+        table = Table(title=f"{kind} (threshold={workload.threshold}) -- {workload.description}", title_justify="left")
+        table.add_column("N", justify="right", style="cyan")
+        table.add_column("Links", justify="right")
+        table.add_column("Before (ms)", justify="right")
+        table.add_column("After (ms)", justify="right", style="green")
+        table.add_column("Speedup", justify="right", style="bold green")
+        table.add_column("Before (MB)", justify="right")
+        table.add_column("After (MB)", justify="right", style="yellow")
+        table.add_column("Peak RAM", justify="right", style="bold yellow")
+
+        for n in n_values:
+            batch = _make_batch(n, kind, dim)
+            timings = _interleaved_best(
+                _baseline, compute_semantic_links_within_batch, batch, workload.threshold, top_k, repeats
+            )
+            before_mb = _peak_mb(_baseline, batch, workload.threshold, top_k)
+            after_mb = _peak_mb(compute_semantic_links_within_batch, batch, workload.threshold, top_k)
+            table.add_row(
+                str(n),
+                f"{timings.links:,}",
+                f"{timings.before_ms:.2f}",
+                f"{timings.after_ms:.2f}",
+                f"{timings.speedup:.2f}x",
+                f"{before_mb:.2f}",
+                f"{after_mb:.2f}",
+                f"-{(1 - after_mb / before_mb) * 100:.0f}%",
+            )
+        console.print(table)
+        console.print()
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument("--sizes", type=int, nargs="+", default=[200, 500, 1700, 5000])
+    parser.add_argument("--dim", type=int, default=1536, help="embedding dimensions (default: 1536)")
+    parser.add_argument("--top-k", type=int, default=50)
+    parser.add_argument("--repeats", type=int, default=5)
+    parser.add_argument("--workloads", nargs="+", choices=sorted(WORKLOADS), default=sorted(WORKLOADS))
+    args = parser.parse_args()
+    run_benchmark(args.sizes, args.dim, args.top_k, args.repeats, args.workloads)
+
+
+if __name__ == "__main__":
+    main()
@@ -43,6 +43,7 @@ client-coverage-check = "hindsight_dev.client_coverage_check:main"
 perf-test = "benchmarks.perf.system_perf:main"
 token-count-bench = "benchmarks.micro.token_counting:main"
 vector-serialization-bench = "benchmarks.micro.vector_serialization:main"
+semantic-within-batch-bench = "benchmarks.micro.semantic_within_batch:main"
 entity-resolver-bench = "benchmarks.micro.entity_resolver_bench:main"
 find-duplicate-observations = "hindsight_dev.obs_dedup.cli:main"

@@ -0,0 +1,15 @@
+#!/bin/bash
+# Microbenchmark the within-batch semantic link pass on the retain path (#3977).
+#
+# Usage:
+#   ./scripts/benchmarks/run-semantic-within-batch-bench.sh
+#   ./scripts/benchmarks/run-semantic-within-batch-bench.sh --sizes 500 5000 --workloads clustered
+
+set -e
+
+SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
+PROJECT_ROOT="$(cd "$SCRIPT_DIR/../.." && pwd)"
+
+cd "$PROJECT_ROOT/hindsight-dev"
+
+exec uv run semantic-within-batch-bench "$@"