Skip to content

fix: Bound reranker input length and attention memory per batch - #836

Draft
buildingjoshbetter wants to merge 4 commits into
maintenance/remaining-code-only-v2from
maintenance/reranker-token-budget-v1
Draft

buildingjoshbetter wants to merge 4 commits into
maintenance/remaining-code-only-v2from
maintenance/reranker-token-budget-v1

Conversation

@buildingjoshbetter

@buildingjoshbetter buildingjoshbetter commented Oct 14, 2026 •

Copy link
Copy Markdown
Owner

The local reranker and the model server ran the cross-encoder with its full 8,192-token limit and a fixed 64-pair batch. Attention scratch grows with batch × length², so one request with long pairs could need tens of GB. On a Mac, that request first fails with an MPS out-of-memory and is then retried on CPU, where it keeps allocating.

This change:

  • caps reranker input at 1,024 tokens (truncated longest side first);
  • scores pairs longest-first in batches sized so batch × tokens² stays within the cost of 64 pairs of 512 tokens;
  • halves that budget below 6 GiB of available memory and quarters it below 2 GiB;
  • leaves the original single predict call untouched when every default batch already fits, and for models without a tokenizer and an integer max_seq_length (proxies, older or custom encoders).

Real ModernBERT reranker on CPU, 100 synthetic candidates, 16 GiB container, peak cgroup memory before → after:

Candidate length Before After Scores
50 words 0.86 GB 0.69 GB identical
400 words 1.39 GB 1.40 GB identical
1,500 words 3.87 GB 1.10 GB change (1,024-token cap)
6,000 words killed at 16 GiB 1.10 GB before never completed

Validation at f1f7655, stacked on #835:

  • Combined suite: 4,515 tests / 5,048 subtests passed (14 new tests), 19 skips, 3 network cases deselected.
  • Focused suites: 782 / 2,336 and 1,327 / 1,469 passed.
  • Installed wheel smoke: 94 Python source files verified.
  • Separate original gates: 739 native source/inert cases and 276 memory cases passed.
  • Independent review accepted after one reported defect was fixed.

Limitation: a cross-encoder that reports max_seq_length as None is neither capped nor budgeted. The shipped ModernBERT and MiniLM rerankers report integer limits.

Base is maintenance/remaining-code-only-v2 (#835). Draft; no merges or releases.

This branch has not been deployed

No deployments
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.

1 participant