slow-J opened a new pull request, #16790: URL: https://github.com/apache/lucene/pull/16790
Starting as a draft, as I want to make sure this approach is valid first. I was looking into more opportunities for vectorization after we shipped https://github.com/apache/lucene/pull/16631 . `OptimizedScalarQuantizer.scalarQuantize` centers the vector, then runs a coordinate descent to pick the quantization interval. Each descent iteration walks the whole vector twice, once for the loss and once for the next step, one scalar double at a time. Then it rounds each component to its quantized value. I am moving those loops into three `VectorUtilSupport` methods: - `osqCenter`: centering, plus the centroid dot product, min, max and sums - `osqGridStats`: one descent iteration's loss and step sums, in one pass - `osqAssign`: rounding to the final quantized values The Panama versions keep a float running sum per lane, with FMA, and add the lanes in lane order at the end (`reduceLanes(ADD)` adds in a different order before and after C2 compiles it, which made the output depend on JIT warm-up). The scalar versions in `DefaultVectorUtilSupport` do the same work in one pass, summing in float for `osqCenter` and in double for `osqGridStats`. The variance, loss and descent sums are now reordered and partly computed in float, so they are less precise. The output is not identical to before: a slightly different sum can end the descent on a different interval, or move a component on a rounding boundary to the next level. Per-doc OSQ loss gets worse about as often as it gets better. On the first 20K Cohere v3 1024d docs: | bits | docs with any value changed | components changed | max change (levels) | OSQ loss vs before (per doc, >0.1%) | |---:|---:|---:|---:|---| | 1 | 2 (0.010%) | 0.0001% | 1 | 0 worse, 0 better | | 2 | 10 (0.050%) | 0.0004% | 1 | 8 worse, 3 better | | 4 | 49 (0.245%) | 0.0048% | 1 | 27 worse, 27 better | | 7 | 327 (1.635%) | 0.0258% | 3 | 129 worse, 134 better | | 8 | 640 (3.200%) | 0.0371% | 4 | 208 worse, 193 better | Recall is unchanged in the luceneutil runs below. ### Benchmarked via JMH Added a `scalarQuantize` pair to `VectorUtilBenchmark` (4 bits, rotating through 64 random vectors) and ran ``` java --module-path lucene/benchmark-jmh/build/benchmarks \ --module org.apache.lucene.benchmark.jmh \ 'VectorUtilBenchmark.scalarQuantize.*' -p size=384,768,1024,4096 -tu ms ``` Hardware: - EC2 Host: m6g.16xlarge - AWS Graviton2 (Neoverse N1), aarch64, no SVE, preferredBitSize=128; FMA enabled, JDK 25.0.2 `scalarQuantizeVector`: ops/ms, main 3bb5f43990 -> this PR | size | main | this PR | ratio | |---:|---:|---:|---:| | 384 | 40.845 ± 1.306 | 348.335 ± 8.971 | 8.53x | | 768 | 18.729 ± 0.625 | 172.376 ± 5.507 | 9.20x | | 1024 | 14.008 ± 0.908 | 126.265 ± 4.294 | 9.01x | | 4096 | 3.111 ± 0.183 | 30.205 ± 1.167 | 9.71x | `scalarQuantizeScalar` (no `--add-modules jdk.incubator.vector`): ops/ms | size | main | this PR | ratio | |---:|---:|---:|---:| | 384 | 41.614 ± 2.166 | 70.515 ± 3.753 | 1.69x | | 768 | 19.235 ± 0.829 | 33.504 ± 1.473 | 1.74x | | 1024 | 14.245 ± 0.700 | 24.396 ± 1.009 | 1.71x | | 4096 | 3.133 ± 0.115 | 5.571 ± 0.156 | 1.78x | ### Benchmarked via luceneutil `knnPerfTest.py` set up like #16631: Cohere v3 Wikipedia 1024d, `ndoc=1_000_000`, `maxConn=64`, `beamWidthIndex=250`, `quantizeBits=4`, `forceMerge`, rerank off + at 32/8/4 bits, `numIndexThreads=8`, `numMergeThread=8`, `numMergeWorker=24`. First with search and stats (topK 100, fanout 100, overSample 1/2/5, 10K queries), then index only. 3 reps each of baseline + candidate. Build seconds (index + merge wait + force merge), mean of 3 reps, main 3bb5f43990 -> this PR | rerank | search + stats: total | Δ total | Δ index phase | index-only: total | Δ total | Δ index phase | |---|---:|---:|---:|---:|---:|---:| | 32 bits | 1891.0 -> 1764.7 | -6.7% | -3.4% | 1896.4 -> 1777.6 | -6.2% | -4.3% | | 8 bits | 963.1 -> 742.0 | -22.8% | -11.2% | 925.9 -> 765.8 | -17.3% | -9.4% | | 4 bits | 963.1 -> 752.3 | -21.4% | -11.2% | 954.1 -> 797.4 | -16.3% | -10.8% | | no | 870.0 -> 798.5 | -8.1% | -5.8% | 855.3 -> 767.8 | -9.6% | -5.5% | | **12 builds** | | **-14.7%** (t -5.2) | **-7.9%** (t -7.2) | | **-12.4%** (t -5.4) | **-7.5%** (t -8.8) | The index phase is faster in all 24 pairs. The totals swing more between reps (-6% to -18% per rep) because faster indexing changes which merges TieredMergePolicy picks, which moves time between merge wait and force merge. Recall and latency, mean of 3 reps, main -> this PR | rerank | overSample | recall (same) | latency main (ms) | latency this PR (ms) | |---|---|---|---:|---:| | 32 bits | 1/2/5 | 0.914/0.995/0.999 | 3.85/6.41/12.27 | 3.81/6.35/12.21 | | 8 bits | 1/2/5 | 0.914/0.981/0.984 | 3.93/6.59/12.43 | 3.83/6.56/12.25 | | 4 bits | 1/2/5 | 0.914/0.918/0.920 | 3.87/6.52/12.11 | 3.74/6.26/11.82 | | no | 1/2/5 | 0.914/0.918/0.920 | 3.73/6.35/11.74 | 3.69/6.17/11.69 | Latency is 1.7% lower over the 36 paired runs (t -5.2). -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected] --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
