Skip to content

Add cupy backend - #188

Closed
lm-cyber wants to merge 4 commits into
xhluca:mainfrom
lm-cyber:add_cupy
Closed

Add cupy backend#188
lm-cyber wants to merge 4 commits into
xhluca:mainfrom
lm-cyber:add_cupy

Conversation

@lm-cyber

@lm-cyber lm-cyber commented Jun 8, 2026

Copy link
Copy Markdown

Summary

Add an explicit, opt-in CuPy backend for BM25 retrieval and top-k selection.

This PR introduces backend="cupy" for GPU retrieval and backend_selection="cupy" for GPU top-k selection, while keeping backend="auto" unchanged: it still selects Numba when available and otherwise falls
back to NumPy.

What changed

  • Added bm25s.cupy with CuPy top-k selection and retrieval utilities.
  • Added CuPy retrieval support to BM25.retrieve(...).
  • Added support for backend_selection="cupy" in the existing selection API.
  • Added sparse/candidate CuPy retrieval paths to avoid scanning all documents when sparse query candidates are enough.
  • Added a synthetic benchmark for comparing Numba, dense CuPy, and sparse CuPy retrieval paths.
  • Added CuPy tests for:
    • missing dependency behavior
    • top-k selection
    • backend retrieval
    • mmap retrieval
    • weight masks
    • sparse batch retrieval
    • candidate/SpMM batch retrieval
  • Updated README and test docs with CuPy installation and test instructions.
  • Fixed token-id lookup so vocabulary id 0 is preserved instead of being dropped as falsy.

Notes

CuPy is intentionally not selected automatically because users need a CuPy package matching their CUDA runtime, for example cupy-cuda12x. Users must explicitly choose backend="cupy" or
backend_selection="cupy".

Testing

Ran locally:

uv run --with numpy --with scipy --with tqdm --with orjson --with PyStemmer python -m unittest tests/core/test_init_utils.py
uv run --with numpy --with scipy --with tqdm --with orjson --with PyStemmer python -m unittest discover tests/cupy
git diff origin/main...HEAD --check

Results:

- tests/core/test_init_utils.py: 14 tests passed
- tests/cupy: 23 tests passed, 20 skipped locally because CuPy/CUDA runtime is not available on this machine
- diff check passed

@lm-cyber lm-cyber changed the title Add cupy Add cupy backend Jun 8, 2026
@xhluca

xhluca commented Jun 8, 2026

Copy link
Copy Markdown
Owner

Thanks so much for putting this together! This looks like a nice implementation, and I can see it being valuable for GPU users.

That said, I'm hesitant to take on cupy as a new dependency for the core library. I'm not very familiar with it myself, and since I expect to keep adding features down the line, I'm worried about introducing something I might inadvertently break later.

Would you be open to releasing this as a separate fork instead, something like bm25s-cupy? That way it can evolve independently and GPU users can still benefit from your work. Happy to link to it from the main repo so people can find it easily. Thanks again for the contribution!

@xhluca xhluca closed this Jun 8, 2026
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.

2 participants