Skip to content

feat: add tiled exact search primitives and optional GPU scoring support - #19

Open
AxelRolov wants to merge 3 commits into
AxelRolov:mainfrom
0rkhann:feat/tiled-gpu-search-primitives
Open

feat: add tiled exact search primitives and optional GPU scoring support#19
AxelRolov wants to merge 3 commits into
AxelRolov:mainfrom
0rkhann:feat/tiled-gpu-search-primitives

Conversation

@AxelRolov

Copy link
Copy Markdown
Owner

Summary

  • add lazy package exports for scoring and optimization to avoid eager import cycles
  • add scoring_tiled.py with tiled exact-search utilities for squared norms, batched Tanimoto, and streamed top-k selection
  • add an optional CuPy-based GPU kernel bridge plus a gpu extra so CPU-only installs continue to work unchanged

Details

This PR prepares cdr_bench to act as a reusable compute backend for large-scale exact search workloads.

The CPU layer adds:

  • compute_fp_squared_norms
  • tanimoto_similarity_matrix_batch
  • topk_1d / topk_matrix
  • exact_tanimoto_topk_from_block
  • streaming_exact_topk
  • resolve_tanimoto_backend

The GPU layer adds:

  • gpu_kernels.py with a safe CuPy import wrapper
  • CuPyKernelBridge with a RawModule fused top-k kernel path and a CuPy fallback path
  • optional dependency wiring under [project.optional-dependencies].gpu

Validation

  • compiled scoring_tiled.py and gpu_kernels.py
  • import smoke for from cdr_bench.scoring import compute_fp_squared_norms, streaming_exact_topk
  • verified the GPU guard path on a CPU environment returns has_gpu_support() == False instead of breaking imports

Notes

The GPU bridge is intentionally optional and isolated. A CPU-only uv sync should continue to work without cupy installed, while GPU users can opt in with uv sync --extra gpu.

@AxelRolov AxelRolov self-assigned this Apr 14, 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