提出Flash-MaxSim优化检索模型内存占用，大幅提升大规模检索效率。
FLASH-MAXSIM: IO-Aware Fused Kernels for Late-Interaction Retrieval
- 设计无中间张量的融合内核，避免存储完整相似度矩阵
- 推理内存减少9倍，训练内存降低100倍，支持更大候选集和批大小
- 兼容现有框架，支持FP32精确计算与INT8低存储版本
晚期交互检索（如ColBERT、ColPali）通过MaxSim操作计算查询与文档的分词级相似度。标准PyTorch实现需生成完整的查询词×文档词相似度张量，再进行压缩，导致内存开销巨大。在ColPali规模下，该张量是流水线中最大内存占用（如10,000文档时达21 GB FP16）。本文提出Flash-MaxSim（FM），一种面向输入输出的融合GPU内核，无需显式构造该张量即可计算相同结果，并将原理扩展至训练反向传播。在A100上，推理内存最高减少9倍，训练内存下降两个数量级，使单卡可处理此前无法达到的候选集与对比学习批大小。内核为即插即用，精度与FP32参考结果一致（评估中nDCG@10误差小于5e-4，BEIR与REAL-MM-RAG数据集），另提供INT8路径以一半索引存储换取高保真度。代码已开源。
原文摘要 · Abstract (English)
Late-interaction retrieval (ColBERT, ColPali) scores a query against a document via the MaxSim operator. The standard PyTorch implementation materialises the full query-token x 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 this cuts inference memory up to 9x and training memory by two orders of magnitude, unlocking candidate sets and contrastive batch sizes a single GPU could not previously reach. The kernel is a drop-in replacement, exact up to floating-point evaluation order under its stated FP32-accumulation protocol: rankings match the FP32 reference within 5e-4 of nDCG@10 on BEIR and REAL-MM-RAG. A separate INT8 path trades exactness for halved index storage at high fidelity. Released open-source.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。