TileMaxSim：分维度分块与融合乘积量化加速GPU检索评分

ColBERT 等多向量检索模型通过细粒度 词元级 MaxSim 评分实现了最先进的准确性，但现有的 GPU 实现并未充分利用大多数硬件性能。我们对现代 GPU 上的 MaxSim 进行了屋顶线分析，并发现了严重的带宽差距：简单的实现仅达到峰值 HBM 带宽的 5-18%，因为它们具体化了 Nq x Nd 相似性矩阵，将内存流量浪费在消耗一次并丢弃的数据上。我们提出了 TileMaxSim，一个 IO 感知的 Triton 内核系列，它通过以下方式缩小了这一差距：(1) 多查询 SRAM 平铺，通过共享内存流式传输文档嵌入，同时在寄存器中累积每个查询词元最大值，从 HBM 读取每个嵌入一次； (2) 维度平铺，将嵌入维度划分为 128 宽的块，从而能够对溢出共享内存的 d > 128 嵌入进行评分； (3) 通过共享内存查找表融合乘积量化评分，将 HBM​​ I/O 削减高达约 31 倍。在 NVIDIA H100 GPU 上，TileMaxSim 达到峰值 HBM 带宽的 80.2%，并得分 82M 文档/秒（真实 MS MARCO 通道上为 71.6M/s），比基于循环的评分加速 220 倍，比融合 PyTorch 加速 6.5 倍，比 torch.compile 加速 6.6-8.5 倍，是同一节点上 WARP CPU 引擎评分吞吐量的 469 倍。 TileMaxSim 保持精确的检索质量：在 MS MARCO 和三个 BEIR 基准上，排名与参考 MaxSim 匹配。作为 ColBERTv2/PLAID 的直接替代品，它将 100K 候选者的评分延迟从 268 毫秒缩短至 1.2 毫秒（端到端延迟降低 98%）。我们进一步展示了从 100K 到 500K 文档的恒定吞吐量、数据并行多 GPU 分片、跨维度 64-768 的鲁棒性以及 FP16/BF16/FP32 支持。并行工作独立开发了IO感知融合MaxSim内核；我们在 d > 128 的维度平铺和融合产品量化评分方面有所不同。