Skip to content

Precompute cosine normalization for indexed pairs - #138

Draft
afermg wants to merge 2 commits into
mainfrom
perf/cosine-pairs
Draft

Precompute cosine normalization for indexed pairs#138
afermg wants to merge 2 commits into
mainfrom
perf/cosine-pairs

Conversation

@afermg

@afermg afermg commented Aug 28, 2026

Copy link
Copy Markdown
Collaborator

Summary

For the exact distance="cosine" AP paths:

  • gather and normalize referenced feature rows once
  • reuse the prepared rows for positive and negative indexed pair dot products
  • avoid repeatedly normalizing copied pair batches
  • preserve exact multiply-then-sum arithmetic, float32 output, tie ordering, and generic/custom-distance fallback behavior

Only referenced rows are normalized, preserving prior behavior for unpaired zero/nonfinite rows. C-, Fortran-, and strided inputs are gathered into the same C-order arithmetic layout as the previous batched implementation.

Benchmark

Synthetic 5,000 × 196, 100,000 pairs, four workers:

origin/main This PR
Wall time 98.74 s 74.84 s including one-time preparation
Peak RSS 596 MB 418 MB
Output SHA-256 identical SHA-256 identical

That contended run measured 1.32× and 30% lower RSS. An earlier real JUMP-lite worker sweep on the same implementation measured 1.3–4.7× kernel speedups depending on batch/worker count, with exact checksums. A current real end-to-end run retained the exact AP digest and reduced peak RSS from 2.86 GB to 1.61 GB, although host contention obscured wall-time improvement.

Validation

  • 47 focused compute/AP tests passed
  • exact float32/float64 parity for C/Fortran/strided inputs and ties
  • referenced/unreferenced zero and infinity behavior covered
  • regular and multilabel paths covered
  • empty pairs and generic metric/callable fallback covered
  • Ruff and git diff --check passed
  • independent follow-up review: OK

@afermg
afermg marked this pull request as draft September 2, 2026 17:13
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant