Daniel Ezer
Author directory2026
FLASH-MAXSIM: IO-Aware Fused Kernels for Late-Interaction Retrieval
Roi Pony | Daniel Ezer | Adi Raz Goldfarb | Idan Friedman | Oshri Naparstek | Udi Barzelay
Proceedings of the 2026 Conference on Empirical Methods in Natural Language Processing: Industry Track
Roi Pony | Daniel Ezer | Adi Raz Goldfarb | Idan Friedman | Oshri Naparstek | Udi Barzelay
Proceedings of the 2026 Conference on Empirical Methods in Natural Language Processing: Industry Track
Late-interaction retrieval (ColBERT, ColPali) scores a query against a document via the MaxSim operator. The standard PyTorch implementation materialises the full query-token × document-token similarity tensor only to reduce it away. At ColPali scale this is the single largest tensor in the pipeline (e.g. 21 GB in FP16 for 10K documents) and limits both candidate set size at inference and batch size during contrastive training. We present Flash-MaxSim (FM), an IO-aware fused GPU kernel that computes the same MaxSim scores without ever materialising the tensor, and extends the same principle to the training backward. At ColPali scale on A100, FM reduces inference peak memory by 1.4-2.6× relative to the deployed chunked baseline (4.9-8.9× relative to unchunked eager execution) and reduces MaxSim-operator training memory by two orders of magnitude, enabling exact reranking over larger resident candidate pools and contrastive batch sizes that vanilla autograd cannot fit on a single GPU. The kernel is a drop-in replacement, exact up to floating-point evaluation order under its stated FP32-accumulation protocol: nDCG@10 differs from the FP32 reference by at most 5×10-4 on BEIR and REAL-MM-RAG. A separate INT8 path trades exactness for halved index storage at high fidelity. Released open-source at https://github.com/roipony/flash-maxsim.