[MOD-17526] Compute SQ8-FP32 L2 via direct residual accumulation, not cancellation - #1028
Open
dor-forer wants to merge 10 commits into
Open
[MOD-17526] Compute SQ8-FP32 L2 via direct residual accumulation, not cancellation#1028dor-forer wants to merge 10 commits into
dor-forer wants to merge 10 commits into
Conversation
… cancellation SQ8_FP32_L2Sqr and its SIMD variants computed L2^2 via the algebraic identity ||x||^2 + ||y||^2 - 2*IP(x, y), reading precomputed sums from blob metadata. This catastrophically cancels in FP32 when x and y share a large common offset relative to their spread (e.g. x=[100000,100008], y=[100000,100000]: true L2^2=64, the old kernel could return 0/negative/garbage). Replace it with direct residual accumulation: dequantize each stored byte, subtract the query value, square, and accumulate. The scalar and SSE4/AVX2 paths no longer read y_sum/x_sum_sq/y_sum_sq. The FMA-capable paths (AVX2_FMA, AVX512, NEON, SVE) fuse the dequantize-and-subtract into a single FMA (diff = fma(delta, q, min - y)) rather than computing the subtract separately, since that fusion is the difference between a 10-23% and a 60-80% slowdown on ARM. SQ8-FP16 L2, SQ8-SQ8 L2, and all IP/Cosine kernels are unchanged.
Covers the gap that let the original bug ship: translation-invariance (L2(x,y) == L2(x+C,y+C)) at realistic dims (128-1024), the ticket's literal large-offset repro, and a regular-vector sanity check -- all checked against an independent double-precision reference rather than another kernel variant sharing the same formula.
…ounding Scalar diff was computed as (min_val + delta*q) - y, left-to-right: the dequantized value gets rounded to FP32 at the large offset's precision before y is subtracted, discarding the residual and reintroducing the cancellation this kernel exists to avoid (only the SIMD kernels used the correct delta*q + (min_val - y) order). Affects x86 dims <8 and any scalar-fallback path. Also fixes the ticket regression test's operands (storage and query were swapped, so the quantized byte was always 0 and delta*q was never exercised), tightens the translation-invariance tolerance now that the dominant error source is gone (min-y is Sterbenz-exact, not a lossy large-offset subtraction as the prior comment claimed), and adds a dedicated regression pinning the exact associativity bug.
![cursor[bot]](https://avatars.githubusercontent.com/in/1210556?s=60&v=4){kind=small}
Codecov Report✅ All modified and coverable lines are covered by tests. Additional details and impacted files@@ Coverage Diff @@
## main #1028 +/- ##
==========================================
+ Coverage 97.43% 97.45% +0.02%
==========================================
Files 141 141
Lines 8686 8843 +157
==========================================
+ Hits 8463 8618 +155
- Misses 223 225 +2 ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
…contract
The test helper calls every ISA kernel directly, bypassing the chooser,
and one case runs at dim=2. On x86 the chooser floors SIMD at dim >= 8
("Optimizations assume at least 8 elements"), and the AVX2 residual path
reads a full 32-byte vector through my_mm256_maskz_loadu_ps, which is an
unconditional _mm256_loadu_ps -- a 16-byte overread of a dim=2 query blob
that ASAN reports. Gate the x86 tiers on the same floor production uses;
aarch64 has no such floor (its tail uses per-lane loads) so NEON/SVE stay
ungated to keep matching production.
Also: spell out the missing `+ offset` on the SVE partial-chunk query load
(zero today only because that block runs first), extend the
translation-invariance dims to 129/145/513/527 so the large-offset case
exercises the residual tails and not just aligned main loops, note that
-ffast-math would reassociate the fix away, and apply clang-format.
…ey don't need
Four independent costs, each measured with the repo's own bm_spaces_sq8_fp32
(median of 5 reps, cv <= 1%).
NEON: the 4-element step read 8 bytes (vld1_u8) but consumed only the low 4
and advanced by 4, so consecutive steps re-read half of every load. Widening
the main loop to one 16-byte load feeding all four accumulators, with the
per-lane arithmetic and accumulator mapping unchanged. Neoverse-N1:
dim16 10.9->8.59ns, dim64 34.5->25.8ns, dim256 128->95.2ns, dim1024
502->370ns (-21% to -26%). The 4-element helper stays for the tail.
SVE: `chunk` is svcntw(), a runtime value, so `dimension % chunk` and
`/ chunk_size` emitted real udiv (~12-20 cycles, unpipelined) on every call.
They were also redundant -- CHOOSE_SVE_IMPLEMENTATION already divides once at
chooser time and passes the answers down as template params. Restructured to
full vectors first with a compared bound and an svwhilelt-predicated tail,
which also keeps every unpredicated load at a multiple of the vector length
instead of pushing them off a leading partial chunk. Isolated like-for-like
compile of two instantiations: 3 udiv before, 0 after; 0 udiv across all 32
instantiations in both SVE.cpp.o and SVE2.cpp.o. Loop algebra verified
exhaustively off-hardware for chunk in {4,8,16,32,64} x dim in [1,4096]:
every element covered exactly once, no overread.
SSE4: the residual path stored scalars into two 16-byte stack arrays and
immediately reloaded them with _mm_load_ps, which cannot be store-to-load
forwarded from narrower stores. It showed in the residual sweep as a fixed
penalty on every dim where residual % 4 != 0 (125-127ns) versus the dims
that skip the path entirely (116-117ns). Now both operands load at full
width -- in bounds because the kernel is unreachable below dim 8 -- with a
compile-time lane mask zeroing the lanes the main loop will handle. The
sweep is flat at 116-119ns.
AVX2 / AVX2+FMA: my_mm256_reduce_add_ps spills 8 floats and sums them with 7
dependent scalar adds, a fixed cost that dominates at small dims. Added
my_mm256_reduce_add_ps_tree (three in-register steps) alongside it rather
than changing it, since the original has callers across the repo that would
each need their own benchmarking. AVX2+FMA dim16 4.42->3.10ns (-30%), dim64
9.05->7.65ns, dim1024 85.4->81.9ns (-4%); AVX2 dim16 5.25->3.91ns (-26%).
AVX512 keeps _mm512_reduce_add_ps, which is already a tree.
Verified: 595/595 SQ8 tests and 595/595 under ASAN on x86 (AVX512/AVX2+FMA/
AVX2/SSE4), 593/593 on aarch64 (NEON). Left alone deliberately, per
measurements: AVX512 stays at 2 accumulators, the scalar path stays at 1
(it only ever sees dims 1-7), and the residual arithmetic itself is unchanged.
There was a problem hiding this comment.
Cursor Bugbot has reviewed your changes using high effort and found 1 potential issue.
❌ Bugbot Autofix is OFF. To automatically fix reported issues with cloud agents, have a team admin enable autofix in the Cursor dashboard.
Reviewed by Cursor Bugbot for commit 47c4f9d. Configure here.
The restructured tail accumulated with svmla_f32_z, whose zeroing form clears every lane outside the predicate. sum3 already holds full-vector partial sums by then -- the main loop's fourth step writes it -- so any dimension with a partial chunk and at least one main-loop iteration (dim >= 4*chunk && dim % chunk != 0) silently dropped part of the distance. The pre-restructure kernel used _z safely here only because the partial block ran first, against an all-zero sum0; moving it after the main loop removed that precondition. Use the merging form instead. The temporaries above stay _z since they are fresh values rather than accumulators. Caught by Cursor Bugbot and then reproduced on Graviton4 (Neoverse-V2, 128-bit VL so chunk=4): the parameterized SQ8_FP32_L2Sqr sweep failed on exactly the dims where dim % 4 != 0 and passed on the rest. With the fix, 593/593 SQ8 tests pass on that host -- the first run of these kernels on hardware that actually has SVE. My off-hardware coverage check did not catch this: it verified which elements each step touches, not what the predicate does to inactive lanes.
The restructure is not a uniform win and the numbers should be in the file rather than only in the PR: on Graviton4 (128-bit VL, median of 9, cv <= 0.33%) SVE2 improves everywhere while plain SVE improves only where the dimension divides evenly by the vector length and costs ~8% where it does not. Kept because SVE2 is preferred when both tiers are present and real embedding dims land on the faster path. Also records the tail variant that measured worse, and why a bitmask cannot replace the modulo (SVE vector length is a multiple of 128 bits, not necessarily a power of two).
…tion Comments on the changed lines drop from 226 to 128. Removed the narration that restates the intrinsic on the next line, and moved the benchmark tables and the rejected-alternative notes out of the sources -- they are already in the commit messages and the PR, which is where they belong. What stays is the reasoning a reader cannot recover from the code: why the operand order is load-bearing, why the SVE tail merges instead of zeroing, why the wide residual load is in bounds, and why the SVE loop shape differs from its siblings. Also switch the SVE main loop from while to for, matching NEON and the sibling SVE kernels. The step helper no longer mutates the caller's offset: it takes one by value, so the cursor advances in the loop header and each unrolled step names the sub-offset it reads. Same instruction sequence -- 0 udiv in both SVE.cpp.o and SVE2.cpp.o, and Graviton4 medians are unchanged (SVE 1024 153ns / 513 84.6ns, SVE2 1024 141ns / 513 79.4ns). 595/595 SQ8 on x86, 593/593 on Graviton4.
The three x86 kernels drove their main loop off a pEnd1 sentinel while their NEON and SVE counterparts count chunks, so the family this PR touches read three different ways. They now all count chunks. Scoped to SQ8-FP32 L2: the sentinel is still the shape in ~32 other x86 kernels, and churning those is unrelated to this fix. Costs nothing: the residual blocks consume dimension % 32, leaving exactly 32 * floor(dimension / 32), and `dimension / 32` divides by a literal, so it is a shift rather than the runtime udiv the SVE version had to avoid. Accumulator counts deliberately unchanged -- AVX2 and AVX2+FMA already use four, and AVX512 stays at two because that was measured faster than four (95.0ns vs 95.9ns at dim 1536). 595/595 SQ8 tests, and 595/595 under ASAN, on x86. Benchmarks land within run-to-run noise, as expected for an unchanged instruction sequence.
…p shape" This reverts commit b57df30.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
SQ8_FP32_L2Sqrand its SIMD variants (SSE4/AVX2/AVX2-FMA/AVX512/NEON/SVE/SVE2) computed L2² via||x||²+||y||²-2*IP(x,y). That identity is exact in real arithmetic but catastrophically cancels in FP32 when x and y share a large common offset relative to their spread: for x=[100000,100008], y=[100000,100000] the true L2² is 64 and the old kernel could return 0, a negative number, or unbounded garbage. Because these kernels serve HNSW graph construction as well as query evaluation, wrong distances select wrong neighbours and the damage is not recoverable by re-querying.Replaced with direct residual accumulation across every ISA variant, keeping the existing loop/residual structure.
Correctness
diff = delta*q + (min_val - y). The operand order is load-bearing:min_valandyare close in magnitude so their subtraction is exact (Sterbenz), and the smalldelta*qcorrection is added afterward. Computing(min_val + delta*q) - yinstead rounds the dequantized value at the large offset's precision and silently reinstates the cancellation — the scalar path originally did exactly that and was fixed during review.-ffast-math/-Ofastwould reassociate this back into the bug. The repo's-O3builds are safe; there's a comment on the kernel saying so.y_sum/y_sum_sqfrom the L2 query blob is deferred:query_metadata_countis keyed by metric, not datatype, so it would change the FP16 L2 layout while FP16 still uses the identity.Tests
New regression tests, all checked against an independent double-precision reference rather than another kernel sharing the same formula (pre-fix, scalar and SIMD were wrong together, so kernel-vs-kernel proved nothing):
L2(x,y) == L2(x+C,y+C)for large C, at dims 128/129/145/512/513/527/768 so the cancellation case also runs through the residual/masked-lane tails.[100000,100008]vs[100000,100000]case.Performance
Kernels got faster, not slower, measured with the repo's own
bm_spaces_sq8_fp32(median of 5-9 reps, cv ≤ 1%):udiv, VL-aligned loads)The SVE loop restructure is not a uniform win and the regression is stated rather than hidden: on the plain SVE tier, dims that leave a tail are ~8% slower (dim 513: 78.3 → 84.6 ns) while dims that divide evenly by the vector length gain. SVE2 improves at both. Kept because SVE2 is the tier the chooser prefers when both are present, and because real embedding dims (128/256/384/512/768/1024/1536) divide evenly at 128- and 256-bit VL. The kernel comment records why the shape differs and points at the ticket; the numbers and the rejected tail variant are in the commit messages.
Left alone deliberately, per measurement: AVX512 stays at 2 accumulators and keeps
_mm512_reduce_add_ps(already a tree); the scalar path stays at 1 accumulator since it only ever sees dims 1-7.Test plan
SQ8suite: 595/595 ondorer-intel(Xeon Platinum 8375C — AVX512/VNNI, AVX2+FMA, AVX2, SSE4)SQ8suite under ASAN: 595/595 on the same hostSQ8suite: 593/593 on Neoverse-N1 (NEON)SQ8suite: 593/593 on Graviton4 / Neoverse-V2 — SVE and SVE2 kernels executed for real. An earlier ARM host was Neoverse-N1, whose feature flags carry nosve, so the runtime-gated SVE/SVE2 checks had been skipping silently and the suite passed without ever entering those kernels.chunk ∈ {4,8,16,32,64} × dim ∈ [1,4096], every element touched exactly once, no overread; plus disassembly confirmingudivcount 3 → 0 (and 0 across all 32 instantiations inSVE.cpp.o/SVE2.cpp.o)Running on real SVE hardware immediately earned its keep: it reproduced a bug Cursor Bugbot flagged in the restructured tail (
svmla_f32_zzeroing live accumulator lanes), failing on exactly the dims wheredim % chunk != 0. Fixed in fc7995d.🤖 Generated with Claude Code
Note
High Risk
Changes core vector distance math used in index build and search; incorrect L2 would corrupt graph topology, though the fix is narrowly scoped to SQ8-FP32 L2 and is heavily regression-tested.
Overview
Fixes MOD-17526 by replacing the SQ8–FP32 L2² formula
||x||² + ||y||² − 2·IPwith direct residual accumulationΣ(dequant(xᵢ) − yᵢ)²in the scalarSQ8_FP32_L2Sqrpath and every SIMD implementation (SSE4, AVX2, AVX2+FMA, AVX512, NEON, SVE/SVE2). The old identity could return 0, negative, or garbage when vectors share a large common offset—breaking HNSW neighbor selection.Per-element math uses
diff = delta·q + (min_val − y)(order is load-bearing for FP32; documented as unsafe under-ffast-math). SIMD kernels no longer delegate to inner-product helpers; they dequantize in the hot loop with FMA where available, multi-accumulator loops, and AVX2 uses new in-registermy_mm256_reduce_add_ps_tree.Adds MOD-17526 regression tests against a double-precision reference: translation invariance at large offsets, ticket repro, scalar associativity, and sanity on normal embeddings; tests exercise dispatched and per-ISA chooser paths where CPU features allow.
Reviewed by Cursor Bugbot for commit 63a2309. Bugbot is set up for automated code reviews on this repo. Configure here.