perf(pq): add transposed SIMD encoder with SGEMM fallback

[jerry-024]

What changed

• Replace the primary IVF-PQ blocked SGEMM encoder with a lazily cached transposed codebook and fused direct-L2 argmin, retaining blocked SGEMM only as the hardware fallback.
• Add specialized AVX2/FMA and NEON kernels for dsub=4, plus generic AVX2/FMA and 16-centroid NEON kernels for other supported shapes.
• Use the transposed path for every batch size when nbits=8, ksub=256, every subvector has at least four dimensions, and the CPU provides a fast fused kernel. x86 CPUs without AVX2+FMA use blocked SGEMM for every batch size; 4-bit and smaller-subvector shapes keep the canonical encoder.
• Add ivf.pq-encoding=auto|canonical. auto is the default; canonical reproduces ProductQuantizer::encode_batch on the same CPU and runtime backend.
• Document and test determinism, batch/thread/split invariance, smallest-index tie breaking, non-finite distance semantics, non-uniform chunks, scalar/SIMD agreement, and transposed-cache reuse/invalidation.

The codebook transpose is built lazily once per ProductQuantizer and reused by subsequent add batches. At d=768 the initialized cache retains 768 KiB. Rust API change: the PQ layout fields (d, m, nbits, dsub, ksub, chunk_offsets), codebook and derived caches are private, with read-only getters for the layout and centroids. set_centroids(...) validates the codebook shape, rebuilds norms, refreshes the cached finiteness flag and invalidates the transpose. Caching finiteness avoids rescanning the full codebook on every small-batch fallback call.

Why

The previous path used SGEMM for each subquantizer and 512-row block, then materialized and scanned a 512 × 256 distance matrix twice for expanded-form correction and scalar argmin. In the original 32,768-row Cohere profile, SGEMM took 46 ms, while correction plus argmin raised the complete path to 333 ms. The transposed kernel accumulates distances and tracks minima without writing the distance matrix.

Encoding behavior

The automatic backend is selected only from shape and CPU features, never from batch size. On x86 with AVX2+FMA and on AArch64, it computes direct squared L2 distances with the transposed encoder. Equal distances choose the smallest centroid index, NaN distances do not update the minimum, infinite distances lose, and code 0 is returned when no distance is below f32::MAX.

Other CPUs use blocked SGEMM for finite 8-bit codebooks with dsub>=4; non-finite codebooks on that hardware and unsupported shapes use the canonical encoder. Both fallback paths use expanded-form arithmetic (|q|² + |c|² - 2q·c). Codes can differ across backends at floating-point ties, and NaN centroids or high-dynamic-range finite inputs can produce different behavior. Set:

ivf.pq-encoding=canonical

when the canonical per-vector encoder must be reproduced on the same CPU and runtime backend. Canonical mode also uses CPU-dispatched norm reductions, so neither mode promises byte-identical codes across CPU feature sets. This setting affects builds only; the index format and query path are unchanged.

PQ encoding inside index_add — Cohere 10M

Measured on 2026-09-08: baseline 27af94dbd120da97a457d2b196b47e23912dbaa5 versus PR 7ff831bfc4c38a16d339a12adccc502018a978b0.

The comparison directly times the existing PQ encoder call inside IVFPQIndex::add_batch, summed over all 10M vectors. It includes encoder dispatch and initial lazy cache construction. Output-buffer allocation, preprocessing, coarse assignment, list insertion, training, temporary-file I/O and serialization are outside the timer. Both isolated source copies have the same benchmark-only Instant instrumentation; elapsed time is captured before logging. No PQ time is inferred from total build or outer index_add time.

Intel Xeon 6982P-C, 8 cores / 16 logical CPUs, AVX2+FMA, 61 GiB RAM; Rust 1.93.1, release build, Rayon 16 threads. Both binaries use the same content-pinned Paimon source and harness, with separate Cargo target directories. Original source checksums, timer patches, compilation paths and instrumented binary checksums are retained. The archived harness has no resolvable Git HEAD.

Cohere 10M, cosine, d=768, nlist=4096, PQ m=192 (dsub=4, 8-bit), OPQ disabled; 262,144 retained training rows. Three paired index-only rebuilds in AB/BA/AB order reuse the same 10 data files and one index shard. Build concurrency=32, range-read concurrency=64, Parquet row-group parallelism=8, max in-flight bytes=805306368. Direct OSS VPC endpoint, Paimon local cache disabled, OS caches uncontrolled; the same disk-backed temporary storage for both variants.

Every run timed exactly 10,000,000 rows in the same 306 calls: 305 batches of 32,768 rows and one batch of 5,760 rows. All six builds and post-build snapshot checks passed.

Pair	Baseline PQ encoding (s)	PR PQ encoding (s)	Time reduction	Speedup
1	96.742	18.715	80.66%	5.17×
2	96.253	17.768	81.54%	5.42×
3	96.515	18.703	80.62%	5.16×

Median PQ encoding time: 96.515 s → 18.703 s. Median paired time reduction: 80.66%; median paired speedup: 5.17×. Paired statistics are the medians of the three per-pair percentages and ratios above.

This measures native AVX2/FMA PQ encoding on this workload. ARM and no-AVX2/FMA small-batch performance remain unmeasured.

Query acceptance — separate uninstrumented run

The earlier 2026-09-08 acceptance run used the same baseline/PR commits without the timing patch. Queries were not repeated for the PQ timing run. Each of its six builds was followed by 5 warmup queries and 1,000 measured queries: topK=100, nprobe=128, no refine, candidate reader binary for both indexes, batch API size 1, request concurrency 8, batch/shard concurrency 14 and range-read concurrency 64. Readers were recreated per request. Snapshot continuity passed after all six query runs.

Pair	Baseline QPS	PR QPS
1	22.365	22.025
2	22.161	21.931
3	22.082	22.458

Metric (median of 3)	Baseline	PR
Recall@100	0.808990	0.808990
NDCG@100	0.853580113	0.853579531
Query peak RSS (GiB)	1.910	1.895
Index bytes	1953300976	1953300976

Median paired QPS change: -1.04%. Acceptance thresholds were Recall/NDCG delta >= -0.002 and median paired QPS change >= -3%; all passed. Top-100 sets differ for 3/1,000 queries and rankings for 20/1,000; repeated per-query results are identical within each encoder. Paired-query bootstrap (10,000 resamples, seed 20260908) 95% CI: Recall delta [0, 0], NDCG delta [-0.000001525, 0.000000014].

Validation

Previously reported functional validation for this PR (not rerun during this benchmark):

• cargo test --workspace
• cargo clippy --workspace --all-targets -- -D warnings
• x86_64 cross-target compilation
• release-mode PQ tests on the macOS AArch64 CI runner
• no-AVX2/FMA dispatch and SGEMM fallback tests under Rosetta x86_64
• Java API and JNI native validation tests
• Python tests: 28 passed

New benchmark verification:

• Direct PQ encoding timing on 7ff831bfc4c38a16d339a12adccc502018a978b0: three paired 10M-row rebuilds, identical batch coverage, verified dependency paths/source/binary checksums and post-build snapshot checks.
• Separate uninstrumented acceptance on the same commits: three paired builds and 1,000 measured queries per index, with query quality and snapshot continuity checks (results above).

[shyjsarah]

Code Review Summary

Mode: full
Scope: GitHub PR #90 at b40b193a0fbda18826332eac795fee3ed63cf4d0 against main merge-base 8dcabf208c99707eab3938af0bcc40530fdb4cad
Files Changed: 10 files (+827/-80)
Score: 71/100
Recommendation: Request changes
GAN Stats: 7 raw findings → 6 after dedup; Discriminator accepted 3 / challenged 2 / rejected 1; Arbiter included 5 / adjusted 2 / excluded 1.

Critical Issues

None.

Major Issues

1. core/src/ivfpq.rs:319-320 — canonical does not deliver the documented cross-CPU byte stability.
   The mode delegates to ProductQuantizer::encode_batch, whose expanded-L2 implementation uses target-dependent SGEMM and norm reductions. A fixed-query/fixed-codebook release probe produced code 1 on aarch64-apple-darwin and code 0 on x86_64-apple-darwin, contradicting the new documentation. Implement a fixed-order architecture-independent canonical encoder, or narrow the public guarantee and documentation.

2. core/src/pq.rs:398-399, 707-753 — AArch64 dsub != 4 is routed from blocked SGEMM to the scalar transposed kernel.
   AArch64 unconditionally enables transposed encoding, but only dsub == 4 has a NEON kernel; other shapes fall through to score_argmin_scalar. On this ARM host, d=768/m=48/dsub=16, 32,768 rows regressed from 91.62 ms on the parent to 125.33 ms on the PR (+36.8%). Make backend selection shape-aware and keep SGEMM for ARM shapes without a benchmarked fast kernel.

Minor Issues

1. core/src/pq.rs:484-514 — The transposed codebook is allocated and rebuilt for every add batch.
   This rewrites m * max_dsub * 256 floats per call (768 KiB for d=768), including one-row streaming calls. Cache derived state safely or provide a writer-owned immutable/precomputed transpose.

2. core/src/pq.rs:396-413 — The no-AVX2/FMA SGEMM fallback now runs for every batch size.
   Removing the previous small-row threshold preserves split invariance, but measured tiny-batch latency regressed by about 52% at one row and 17% at seven rows. Add a deterministic low-overhead tiny-batch fallback or document/benchmark the tradeoff.

3. .github/workflows/ci.yml — New AArch64-only unsafe kernels have no behavior test in PR CI.
   The ordinary Rust/FFI/JNI behavior jobs run on Ubuntu x86_64; release build/load smoke tests cannot validate NEON numerical equivalence, tie handling, or non-finite semantics. Add an AArch64 core PQ test job and a forced non-AVX2 fallback job.

Validation

• cargo fmt --all -- --check: passed.
• cargo test --workspace: passed (core 488 passed / 1 ignored; integration, FFI, and JNI tests passed).
• cargo clippy --all-targets --workspace -- -D warnings: passed.
• cargo check -p paimon-vindex-core --target x86_64-apple-darwin: passed.
• Cross-target canonical reproduction independently confirmed: aarch64=1, x86_64=0.
• Security review found no actionable issue.

Positive Observations

• The PR adds extensive correctness tests for ties, non-finite values, batch/thread/split invariance, non-uniform chunks, and scalar/SIMD agreement.
• The specialized dsub=4 fast path is well targeted to the documented production shape and shows a substantial x86 AVX2/FMA speedup.
• The fallback behavior and arithmetic differences are documented more thoroughly than typical performance changes, although the canonical portability claim needs correction.

[shyjsarah]

Rechecked: the no-AVX2/FMA tiny-batch regression remains, but its scope is limited and the trade-off is reasonable given the deterministic semantics and gains on the primary workloads. I consider this acceptable and non-blocking.

[JingsongLi]

+1