Fix long-context search OOMs, improve search speed with auto batching, add index auto device placement

[paulomouraj]

Why

When evaluating popular retrieval benchmarks with long context settings (query/doc lengths of 8192 tokens), some challenges appeared:

• Need to manually reduce/tune default search batch sizes to avoid GPU OOM during search step. There is no single batch_size value that fits every query length, corpus, and GPU.
• Even with tuning, some long context and code benchmarks still OOM during search step. Cause? Exact scoring step materializes big tensors (scales with query tokens x doc tokens) and no chunking of them was configurable/available. The solution was then to implement chunking for those tensors, but enabling it as default was slowing down search time for datasets that do not require it.
• Some datasets such as MLDR OOMed at GPU index load, so low_memory=True was needed, but this makes search much slower. The solution was then to manually keep residuals on CPU and only codes on GPU.

All of this tuning overhead justifies implementing auto batching/chunking based on free GPU memory, avoiding possible OOMs with long queries/documents and avoiding slowing down the search step for datasets that do not require it. The same requirements applies for index placement, which is dataset specific and requires placing codes or residuals on CPU/GPU based on available GPU memory. In addition to that, some ops were wasting gpu memory: 3D query tensors may be huge depending on the longest query on the batch, and those could be replaced by packed 2D query tensors; MaxSim scoring creating big tensor copies to mask padded positions, when in-place fill-in was possible. Inside large tensor splitting/chunking, sorting documents before chunking enables minimal padding waste and thus lower gpu memory.

Summary of all changes

1. Memory-budgeted scoring chunks. Scoring batch/chunk sizes are computed automatically for both scoring stages (approximate and exact rerank) from a per-device memory budget. The default budget is 50% of the free VRAM (search_memory_fraction=0.5) at query time (sampled with torch). On CUDA OOM the stage retries with the budget halved. Batching/chunking only kicks in when needed: a single-chunk fast path keeps short queries at full speed, since chunking ops may introduce extra time. If users still want to manually tune it, they can do it forcing a int on batch_size='auto' | <int> that will be used as approximate score batch size and exact score chunk size. Sorting is implemented wrt document lengths when chunking is enabled so that minimal padding is used for tensor ops, reducing gpu memory waste and potentially improving speed.
2. Packed query tensors. Queries are passed as one packed 2D tensor plus per-query lengths instead of a padded 3D tensor - no padding memory waste on GPUs and a single host-to-device transfer. query_lengths is now required at the Rust FFI and the padded-3D path is removed; the public Python API is unchanged (packing happens internally).
3. Adaptive index placement. Exposed as index_gpu_memory='auto' | 'low' | 'medium' | 'high' (replaces low_memory): 'high' keeps codes and residuals on GPU, 'medium' keeps codes on GPU and residuals (the larger tensor) on CPU, 'low' keeps both on CPU. 'auto' (default) picks the highest tier that keeps total device occupancy within index_memory_fraction (default 70% of total VRAM), reserving the rest for the search step.
4. In-place MaxSim masking. The mask fill in MaxSim scoring is applied in place on the freshly materialized score tensor, avoiding an extra full-size GPU allocation.
5. No GPU load in create(). FastPlaid.create() no longer loads the index onto GPU; per-device loads are deferred to the first search, since the object built during creation would be discarded anyway.
6. Per-query search errors now propagate instead of being silently returned as empty results.

End-to-end benchmark

Environment

Component	Version
(vanilla) fast-plaid	1.4.7
(PR) fast-plaid	paulomouraj/fast-plaid faster-eval-long-context @ dd94341
pylate	1.5.0
mteb	2.12.25
torch	2.9.0 (cu128)
transformers	5.3.0
sentence-transformers	5.3.0
accelerate	1.13.0
flash-attn	2.8.3
datasets	4.8.4
Python	3.12.13
GPU	7x NVIDIA H100 80GB per build, driver 580.159.03
Model	lightonai/mLateOn (ColBERT), query/doc length 8192

Both builds compiled from source with maturin (release, LIBTORCH_USE_PYTORCH=1) against the same
torch; identical eval stack. Search params both sides:
top_k=1000, n_full_scores=8192, n_ivf_probe=8. Vanilla fast-plaid uses PyLate (1.5.0) defaults (+low_memory=True for MLDR) while PR fast-plaid uses everything on auto (PR defaults).

Per-task results (PLAID search step)

Task	vanilla search (s)	PR search (s)	speedup	vanilla ms/q	PR ms/q	vanilla max mem (GiB)	vanilla mem avg (GiB)	PR mem max (GiB)	PR mem avg (GiB)
MSMARCO	OOM	1316.8	—	—	188.7	—	—	51.5	51.5
FEVER	OOM	4001.5	—	—	600.3	—	—	51.2	51.1
ClimateFEVER	OOM	1597.9	—	—	1040.9	—	—	51.2	51.1
HotpotQA	OOM	192.5	—	—	26.0	—	—	28.1	28.1
DBPedia	OOM	54.0	—	—	134.9	—	—	27.1	27.1
NQ	OOM	33.1	—	—	9.6	—	—	29.1	28.3
QuoraRetrieval	15.9	16.1	1.0x	1.6	1.6	7.1	5.7	1.1	1.0
CQADupstackRetrieval	OOM (92.6s partial)	127.5	—	11.4	9.7	46.6	30.1	8.1	4.2
Touche2020	6.8	9.8	0.7x	139.7	199.0	70.8	66.9	13.9	13.0
TRECCOVID	OOM	4.1	—	—	81.6	—	—	6.2	5.4
FiQA2018	6.2	5.0	1.2x	9.6	7.7	25.2	20.1	1.6	1.5
SCIDOCS	8.9	6.3	1.4x	8.9	6.3	28.7	22.0	1.5	1.4
ArguAna	36.5	10.0	3.7x	26.0	7.1	66.3	66.3	9.8	7.5
SciFact	3.2	3.2	1.0x	10.7	10.6	2.7	2.7	1.6	1.5
NFCorpus	3.4	3.3	1.0x	10.6	10.2	2.0	2.0	1.6	1.6
MIRACLRetrievalHardNegatives	OOM	103.7	—	—	9.4	—	—	12.4	4.0
MultiLongDocRetrieval	OOM (138.2s partial)	1761.8	—	691.1	275.3	32.2	32.1	56.0	43.6
SyntheticText2SQL	13.0	10.5	1.2x	2.2	1.8	6.3	5.4	1.2	0.7
CodeFeedbackMT	OOM	1317.5	—	—	99.2	—	—	39.5	39.2
CodeFeedbackST	OOM	1049.1	—	—	33.5	—	—	41.0	40.7
COIRCodeSearchNetRetrieval	OOM	619.8	—	—	11.8	—	—	47.9	38.6
CodeSearchNetCCRetrieval	OOM	796.0	—	—	15.1	—	—	40.9	37.3
CodeEditSearchRetrieval	59.5	46.3	1.3x	4.6	3.6	6.9	3.1	1.7	0.4
StackOverflowQA	OOM	36.1	—	—	18.1	—	—	36.3	34.9
CosQA	3.0	2.9	1.1x	6.1	5.8	4.5	3.9	0.6	0.6
AppsRetrieval	OOM	42.8	—	—	11.4	—	—	13.4	11.0
CodeSearchNetRetrieval	OOM	24.9	—	—	4.2	—	—	4.6	1.2
CodeTransOceanContest	8.7	3.5	2.5x	39.6	15.9	55.7	55.7	4.9	4.2
CodeTransOceanDL	6.4	2.8	2.3x	35.4	15.6	51.8	51.8	5.5	5.5
Total (12 tasks complete on both)	172	120	1.4x	—	—	—	—	—	—

Notes:

• MultiLongDocRetrieval: vanilla OOMs after first lang (ar);
• CQADupstackRetrieval: vanilla covers 8 of 12 subtasks before OOM.
• Touche2020 (0.7x) is the single row where PR is slower (6.8s -> 9.8s over 49 queries): planner
  overhead on a tiny query set, but PR uses 13.9 GiB while vanilla peaked at 70.8 GiB.
• Some OOMs on vanilla fast-plaid can be avoided by setting low_memory=True (index kept on
  CPU, much slower search) and/or lowering the query batch size (PyLate's retriever passes
  batches of 50 queries into search(); smaller batches shrink every scoring tensor
  proportionally, at a large speed cost). Others cannot: on MLDR the OOM occurs in the
  exact-scoring step with low_memory=True already set, and that tensor has a per-query
  floor of ~34 GiB (n_full_scores=8192 candidates × 8192-token docs × 128 dims), no batch
  size fixes it, and lowering n_full_scores changes retrieval results.