Repository navigation
Fix MaxSim memory use for long query batches - #4153
Merged
tomaarsen merged 3 commits intoOct 13, 2026
Merged
Conversation
Member
|
Hello! I've pushed some changes to simplify the implementation, tests, and docstrings. The mixed-dtype compatibility logic added complexity and slowed ordinary inputs, so I removed it while keeping the query chunking fix. In the #4151 reproduction, scoring was about 12x faster and peak additional CUDA memory fell from 871 MiB to 17 MiB, with identical scores. Most standard cases were unchanged, though the small control case added about 0.15 ms. I think I'll merge this when the tests pass.
|
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Fixes #4151.
A long query could push late-interaction retrieval batches over
chunk_elements. MaxSim padded every query to the longest one and always scored at least one document, leaving the scoring intermediate unbounded.This chunks queries and documents separately, counting both padded inputs and the 4D intermediate against the budget. Query chunks are sized against the widest document. Documents are processed again for each query chunk, and a single pair can still exceed the budget. With 200 16-token queries, one 2,000-token query and four 1,000-token documents (dim 128, fp16, budget 10M), extra peak allocated CUDA memory fell from 869.7 MiB to 18.3 MiB, with identical scores.
Shared masks and padding dtypes behave as before.
Added budget, gradient, mask and dtype regressions.
.venv/bin/python -m pytest -q tests/util/test_similarity.pypasses. The offline full-suite run had existing failures and missing model fixtures.Assisted by Claude/Codex.