[htripathi12]

A batch of 32 requested BGE-M3 query–passage pairs currently computes 1,024 combinations and returns only the diagonal. Compute dense and sparse dot products for corresponding rows and apply ColBERT MaxSim only to those pairs. Process passage tokens in tiles of 512 to bound the temporary similarity allocation for long passages.

Fixes #1624. Related to #1625, which also implements corresponding-pair scoring. This implementation adds tiled MaxSim and public compute_score regression coverage. The two commits separate pairwise scoring from tiling for review.

Behavior and tests

Retain temperature and query-length normalization, output ordering, default/custom score weights, and zero-padded passage positions in MaxSim. The shared training scoring methods retain their existing implementation.

Four public-API regression cases compare all five score outputs against the existing all-pairs scorers. They cover a single pair, batches with mixed lengths, a short final batch, default/custom weights, and negative similarities with zero padding across the 512-token tile boundary.

Measured results

Windows, RTX 5060 8 GB, Python 3.13.16, PyTorch 2.11.0+cu128, Transformers 4.57.6, TF32 disabled. Full-model runs used BAAI/bge-m3 revision 5617a9f61b028005a4858fdac845db406aefb181 with FP16 model weights. Its ColBERT output vectors were FP32.

Workload	Upstream	Pairwise only	Pairwise + tiles
Synthetic ColBERT, B=32, Lq=128, Lp=2048, D=1024	74.25 ms / 1025.50 MiB	3.55 ms / 32.02 MiB	3.70 ms / 8.05 MiB
Synthetic ColBERT, B=32, Lq=128, Lp=8192, D=1024	Skipped by memory guard	12.62 ms / 128.02 MiB	13.46 ms / 8.05 MiB
Full API, 128 SciFact pairs, batch 32, passage cap 512	1606.25 ms / 486.99 MiB	1600.30 ms / 486.99 MiB	1589.48 ms / 486.99 MiB

Synthetic numbers are warmed medians over five measured runs and temporary peak PyTorch allocation above resident input tensors. Full-API numbers are warmed medians over three measured runs and peak allocation above the resident model. Full-API measurements include tokenization and encoding. Timing differences in the full API are small and are not evidence of a substantial full-model speedup.

For a synthetic stress case with B=96, Lq=1024, Lp=8192, D=1024 under a 5705.4 MiB PyTorch CUDA allocation cap, untiled pairwise scoring raised OOM and tiled scoring completed in 300.48 ms with 193.12 MiB temporary allocation. The all-pairs baseline was skipped by a memory guard. This capacity result covers scoring on resident synthetic embeddings.

On 128 SciFact pairs arranged as 16 queries with eight candidates each, all five scores matched upstream within tolerance and all 16 ensemble rankings matched at batches 1, 8, and 32. Candidate selection used a positive passage plus deterministic lexical-overlap candidates.

A separate full-model run concatenated SciFact abstracts into passages truncated to exactly 8192 tokens. With queries capped at 128 tokens and the same 5705.4 MiB allocation cap, upstream and pairwise-only both completed batch 20 and raised OOM at batch 21. Their total peak allocations were both 4631.51 MiB. Median batch-20 latency was 11087.48 ms upstream and 11000.50 ms pairwise-only. All five outputs matched with maximum absolute difference 0.0001221.

Tiling adds loop/launch overhead. In the tested B=1, Lp=8192 scoring case, tiling took 2.23 ms versus 0.37 ms for untiled pairwise scoring. Its benefit is bounded temporary memory.

Reproduce regression checks

Run in separate environments using CUDA-enabled PyTorch 2.11.0 and Python 3.13.

python -m pip install -e . pytest "transformers==4.57.6"
python -m pytest -q tests/test_m3_pairwise_scores.py

python -m pip install -e . pytest "transformers==5.22.2"
python -m pytest -q tests/test_m3_pairwise_scores.py tests/test_imports_v5.py

Fresh checks on the submitted tiled branch passed all four parity cases on Transformers 4.57.6 and all nine parity/import cases on Transformers 5.22.2.

The earlier targeted compatibility selection completed with 16 passed and one Transformers-5-only skip on Transformers 4.57.6. The full pretrained-model repository suite and multi-GPU execution were not run.

AI assistance

OpenAI Codex assisted with implementation, regression tests, GPU validation, and this description.