TileMaxSim: IO-Aware GPU MaxSim Scoring with Dimension Tiling and Fused Product Quantization
针对 ColBERT 多向量检索的 token 级 MaxSim 打分,论文先做 roofline 分析,指出朴素实现因要物化 Nq×Nd 相似度矩阵只用到 HBM 峰值带宽的 5–18%。TileMaxSim 用 IO-aware Triton 内核解决:多 query 的 SRAM tiling 让文档 embedding 只读一次、维度 tiling 支持 d>128、融合乘积量化查表把 HBM I/O 砍掉约 31 倍。H100 上达到峰值带宽 80.2%、每秒打分 8200 万文档,作为 ColBERTv2/PLAID 的 drop-in 替换把 10 万候选打分延迟从 268ms 降到 1.2ms(端到端降 98%)。
INSIGHT多向量/晚交互重排在社区搜索里常因打分太慢上不了线,TileMaxSim 说明瓶颈不在算力而在访存——别物化相似度矩阵、让 embedding 只过一次 HBM。它在 MS MARCO 和三个 BEIR 上排序与参考 MaxSim 完全一致,且 10 万到 50 万文档吞吐恒定、支持多 GPU 分片与 FP16/BF16/FP32,落地门槛低,可直接接到现有 ColBERT 检索栈。