Seismic search optimization by chishui · Pull Request #19 · opensearch-project/neural-sparse-cpp

chishui · 2026-07-15T04:25:56Z

Description

This PR include two search optimization:

Use inverted index for cluster summary similarity score, this increases the query QPS by 70%
Reuse data structure in a same thread, this increases by 1.5%

All benchmark is done on c5.9xlarge ec2 with only 8-cores used for query and the dataset is MSMarco v1 8.8M.

SQ width comparison under exp9 (base_full, 8-thread, k=10, AVX-512)

  ┌─────────┬─────────────────────────┬────────────┐
  │ version │ pre opt                 │ opt now     │
  ├─────────┼─────────────────────────┼────────────┤
  │ float   │ 19,200                  │ 33,083     │
  ├─────────┼─────────────────────────┼────────────┤
  │ int16   │ 20,400                  │ 30,615     │
  ├─────────┼─────────────────────────┼────────────┤
  │ int8    │ 27,900                  │ 42,047     │
  └─────────┴─────────────────────────┴────────────┘

Issues Resolved

List any issues this PR will resolve, e.g. Closes [...].

By submitting this pull request, I confirm that my contribution is made under the terms of the Apache 2.0 license.
For more information on following Developer Certificate of Origin and signing off your commits, please check here.

…namic scheduling Signed-off-by: Liyun Xiu <xiliyun@amazon.com>

zirui-song-18

Great idea to optimize CSR into CSC. Is it possible for you to add some UTs? Or IT to verify the data consistency between existing CSR and CSC.

zirui-song-18 · 2026-07-15T05:54:53Z

+        const term_t t = q_idx[i];
+        const int32_t li = term_lookup_[t];


nit: Is it worth adding a boundary check here for term_id < dim?

zirui-song-18 · 2026-07-15T05:57:29Z

+
+    // Distinct terms present across all summaries, ascending.
+    size_t dim = summaries_->get_dimension();
+    term_lookup_.assign(dim, 0);


It allocates a full dim-sized int32 array (~120 KB at dim≈30k) per posting list. Across ~30k posting lists that's ~3.6 GB. The CSC value/cluster arrays themselves (~2× nnz) are a fine trade, but he dim-sized dense lookup per posting list could be the concern. Is it possible to further optimize here?

change to binary search

zirui-song-18 · 2026-07-15T05:59:58Z

    if (summaries_->num_vectors() == 0) {
        summaries_.reset();
    }
+    build_transpose();


The transpose duplicates summaries_. What still reads the CSR after we get CSC?

remove summaries_

zirui-song-18 · 2026-07-15T06:00:48Z

+    // query term (t, qv), add qv*value to scores[cluster_id]. Avoids the
+    // dense-buffer gather entirely.
+    void build_transpose() const;
+    bool has_transpose() const { return transpose_built_; }


I did not see the caller for this?

zirui-song-18 · 2026-07-15T06:01:36Z

+        const idx_t s = term_ptr_[li - 1];
+        const idx_t e = term_ptr_[li];


nit: I think you mean 'start' and 'end' here? What would be the better practice?

Build a per-posting-list CSC transpose of each cluster's summary vectors at index load, then score summaries by iterating the query's non-zero terms (scatter-add into a small per-list scores array) instead of gathering a dimension-sized dense query buffer once per summary. The old path did an AVX-512 gather of dense[term] for every term of every summary in each scanned list -- a scattered read over a ~120 KB buffer that overflows L1/L2 and dominated search time. The transpose keeps the working set to the query's ~49 active terms plus a few-KB L1-resident scores array. Search throughput +72% at k=10 / +27% at k=100 (base_full, 8.84M docs, 8 threads, AVX-512), with byte-identical top-10 results vs the prior version (recall unchanged). Signed-off-by: Liyun Xiu <xiliyun@amazon.com>

chishui · 2026-07-15T06:34:25Z

Great idea to optimize CSR into CSC. Is it possible for you to add some UTs? Or IT to verify the data consistency between existing CSR and CSC.

the SeismicIndex already has UTs which covers the flows

- Replace the retained dimension-sized term_lookup_ table with a binary search over the sorted distinct summary terms (term_ids_). This removes the per-posting-list dim-sized allocation (~120 KB each), so the transpose's footprint is now proportional to the summaries' nnz rather than the vocabulary dimension. - Query terms outside the summaries' term range are now skipped safely via the lower_bound miss (no unchecked dim-sized index). - Remove the unused has_transpose() accessor. - Rename local s/e to start/end in the scoring/build loops. - Add a unit test asserting the transposed (CSC) summary scores exactly match the reference dense-gather over the CSR summaries, across present, absent, and out-of-range query terms. Signed-off-by: Liyun Xiu <xiliyun@amazon.com>

…ze() The term-major transpose was only built in deserialize(), so an index built in-process (via build()/summarize()) or copied left transpose_built_ == false. score_summaries_transposed() then returned all-zero summary scores, so cluster pruning/ordering degenerated and depended on the random k-means assignment -- observable as a flaky SeismicIndexSearch.search_respects_k_limit (doc2 sometimes returned over the higher-scoring doc1). Rebuild the transpose at the end of summarize() (resetting the built flag first) and copy the transpose members in the copy constructor / assignment operator, so every construction path has a consistent transpose. Fixes the flaky test (100/100 now; whole seismic + cluster suite green over repeated runs). Signed-off-by: Liyun Xiu <xiliyun@amazon.com>

…aries_ Extends the term-major (CSC) summary transpose to be element-size-aware (float / uint16 / uint8 with matching accumulation) and routes the scalar-quantized index's summary scoring through it as well, so both index types share the query-driven scoring path. With the transpose now covering every query path, the CSR summaries_ member is removed entirely: InvertedListClusters stores only the transpose, summarize() builds it directly from the transient per-cluster summary, and serialization reads/writes the CSC form. This addresses the reviewer's question about what still reads the CSR after the transpose is built (nothing, now). Note: this changes the on-disk index format (CSC instead of CSR summaries); existing saved .dat files must be rebuilt. Tests that inspected summaries() are rewritten to assert via the public score_summaries_transposed() (a one-hot query recovers a summary value). Full suite green (362 tests), SQ + seismic + cluster suites stable over repeated runs. Signed-off-by: Liyun Xiu <xiliyun@amazon.com>

Read the quantizer width from NSPARSE_SQ_BITS ("8bit" or "16bit", default "8bit") in BM_SeismicSQ_Search and cache each width to its own .dat suffix, so float / int16 / int8 search throughput can be compared without editing the benchmark. Default behavior is unchanged when the variable is unset. Signed-off-by: Liyun Xiu <xiliyun@amazon.com>

chishui · 2026-07-15T09:29:02Z

+    // per-summary gather into the dimension-sized dense buffer. The summaries
+    // hold float values, so the query values are passed as their raw bytes.
+    cluster_invlist.score_summaries_transposed(
+        q_idx, reinterpret_cast<const uint8_t*>(q_val), q_len, score_scratch);


yes, we are using a general data structure which stores bytes, and for different data type, float, int8, int16, we just cast them to byte then cast bytes back to original types.

zirui-song-18

LGTM

The term-major (CSC) summary transpose stored cluster ids as int32, but clusters per posting list are bounded by beta (a small fraction of lambda) and stay far below 2^16 for any workable config. Narrowing csc_cluster_ to uint16 halves that array at no recall cost. On base_full (8.84M SPLADE docs) this shrinks the summary store ~20.8% (float) to ~30.3% (SQ-8bit) and the on-disk index ~12-19%, with search performance retained (SQ-8bit slightly faster from reduced memory traffic) and recall unchanged. This changes the serialized index format, so cached indexes must be rebuilt. The canonical Rust seismic stores this same field and likewise caps the summary count at 2^16; build_transpose asserts the bound holds. Signed-off-by: Liyun Xiu <xiliyun@amazon.com>

The scalar-quantized search path never received the allocation optimizations that SeismicIndex::search got. It allocated a fresh dimension-sized dense buffer and a visited-doc set per query, and re-materialized the query batch into a SparseVectors on every call. Mirror the float index: reuse a per-thread dense buffer (restored to all-zero via a sparse clear over each query's own dims) and visited set across the queries a thread handles, index the pre-quantized batch codes directly instead of building a SparseVectors (kept only for the id-selector exact-match path), and use schedule(dynamic, 64). Behavior-preserving: all 362 tests pass and recall@10 on base_full SQ-8bit is unchanged (0.8327). Search throughput improves ~6% on base_small (k=10 160.0k -> 169.8k QPS, k=100 66.9k -> 70.7k at 16 threads), where per-query allocation is a meaningful fraction of the work; the gain washes out on the full 8.84M corpus where per-doc scoring dominates. Signed-off-by: Liyun Xiu <xiliyun@amazon.com>

chishui requested review from model-collapse, yuye-aws and zirui-song-18 as code owners July 15, 2026 04:25

Optimize SeismicIndex search: per-thread scratch, drop query copy, dy…

7f3f2c8

…namic scheduling Signed-off-by: Liyun Xiu <xiliyun@amazon.com>

chishui force-pushed the seismic-search-alloc-opt branch from 51a0bda to a8eee23 Compare July 15, 2026 04:28

zirui-song-18 reviewed Jul 15, 2026

View reviewed changes

chishui force-pushed the seismic-search-alloc-opt branch from a8eee23 to 44fa2d6 Compare July 15, 2026 06:35

chishui added 4 commits July 15, 2026 06:47

This comment was marked as resolved.

Sign in to view

chishui commented Jul 15, 2026

View reviewed changes

zirui-song-18 approved these changes Jul 15, 2026

View reviewed changes

chishui added 2 commits July 16, 2026 02:28

chishui merged commit 655279e into opensearch-project:main Jul 16, 2026
8 checks passed

		const term_t t = q_idx[i];
		const int32_t li = term_lookup_[t];

		const idx_t s = term_ptr_[li - 1];
		const idx_t e = term_ptr_[li];

Uh oh!

Conversation

chishui commented Jul 15, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Description

Issues Resolved

Uh oh!

zirui-song-18 left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

chishui commented Jul 15, 2026

Uh oh!

This comment was marked as resolved.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

zirui-song-18 left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

chishui commented Jul 15, 2026 •

edited

Loading