Sync with upstream + add in-line filtered late-interaction search

[jakemannix]

Summary

Two changes on this branch:

1. Fast-forward sync of the fork with stanford-futuredata/ColBERT main (6 upstream commits, no divergence).
2. In-line filtered late-interaction search — a new extension that walks filter inverted lists in parallel with IVF cluster lists during candidate generation, instead of pre- or post-filtering.

In-line filtering design

The current ColBERT path is: walk cluster lists → concat → sort → unique → apply filter_fn → score. The new path runs the filter during the cluster-list traversal, so the sort/unique only sees the already-filtered set. For selective filters this is a big reduction in dedup work, candidate-pool memory, and downstream centroid-pruning work.

New modules:

• colbert/search/filter_index.py
  • FilterIndex.from_field_data(...) — builds inverted lists keyed by (field, value), packed in a StridedTensor of int32 PIDs.
  • FilterExpression (Term / And / Or) — compose filters; .materialize(filter_index, mode='auto'|'bitmap'|'sorted') compiles to a runtime filter.
  • MaterializedFilter: BitmapFilter (bool tensor of length num_passages, branch-free / GPU-friendly) and SortedListFilter (sorted-ascending PIDs + searchsorted).
  • Both expose filter_packed(packed_pids, lengths) → (filtered_packed, filtered_lengths) — the primitive invoked during cell-list traversal.
• colbert/search/filtered_search.py
  • filtered_ivf_lookup(ivf, cell_ids, materialized_filter) — vectorised single-pass implementation.
  • filtered_ivf_lookup_per_cell(...) — reference per-cell merge-join (canonical algorithm; mirror layout for a future per-block CUDA kernel).

Wiring (additive, backward-compatible):

• CandidateGeneration.generate_candidates accepts materialized_filter=... and dispatches to the new code path when provided.
• IndexScorer.retrieve / IndexScorer.rank pass it through. The legacy filter_fn callback still works and stacks on top.
• Searcher.search / search_all / dense_search plumbed end-to-end.

GPU acceleration plan:

• Hybrid (current ColBERT, data in RAM): bitmap path = one indexed gather + one scatter_add_ (CUDA-friendly torch ops); sorted-list path = searchsorted + segment-sum.
• Full-VRAM (small index): same filter_packed(packed_pids, lengths) signature lowers directly to a fused CUDA kernel reading IVF + bitmap in one pass. The _per_cell_reference form intentionally maps to one CUDA block per cell with warp-cooperative merge inside.

Test plan

• colbert/tests/inline_filter_test.py — 27 unit tests, CPU-only, no model/index needed:
  • FilterIndex construction (single, multi-valued, missing values, unknown terms)
  • Expression materialisation (Term, And, Or, nested)
  • Auto-mode density heuristic
  • Both materialised filter primitives (BitmapFilter / SortedListFilter) against a brute-force reference
  • filtered_ivf_lookup matches filtered_ivf_lookup_per_cell matches brute-force reference
  • Randomised correctness: 6 seeds, random IVF + random filter + random cell selection — in-line filtered candidate gen ≡ legacy (union → dedup → filter) candidate set
  • End-to-end through CandidateGeneration.generate_candidates with the materialized_filter kwarg
• Benchmark on a real ColBERT index (deferred; needs GPU + actual residual encoder)
• CUDA kernel implementation of filter_packed for the full-VRAM mode (deferred)

Commits

• cc4f3dc (and 5 prior) — upstream sync
• 546f088 — in-line filtered search