优化大模型检索速度,实现每秒8200万次文档评分
TileMaxSim: IO-Aware GPU MaxSim Scoring with Dimension Tiling and Fused Product Quantization
- 用分块共享内存和融合量化减少显存读写,只读一次数据
- 在H100上达到峰值显存带宽的80.2%,提速220倍
- 可直接替换现有模型,适合需要低延迟检索的系统
多向量检索模型如ColBERT通过细粒度的词级最大相似度(MaxSim)评分实现顶尖精度,但现有GPU实现未能充分利用硬件性能。我们对现代GPU上的MaxSim进行屋顶分析,发现严重带宽瓶颈:原始实现仅达峰值显存带宽的5-18%,因需完整生成Nq×Nd相似度矩阵,导致大量一次性数据被重复读取浪费。本文提出TileMaxSim,一系列面向输入输出的Triton内核,通过三方面改进:(1) 多查询共享内存分块,将文档嵌入流式传输至共享内存并寄存器累加每查询词的最大值,每条嵌入仅从高带宽内存(HBM)读取一次;(2) 维度分块,将嵌入维度划分为128位宽的块,支持超过128维的嵌入处理;(3) 融合产品量化评分,利用共享内存查找表,显存访问量减少约31倍。在NVIDIA H100上,TileMaxSim达到80.2%的峰值显存带宽,每秒可评分8200万条文档(真实MS MARCO数据集为7160万/秒),相较循环基评分提升220倍,比融合PyTorch快6.5倍,比torch.compile快6.6-8.5倍,比同节点上WARP的CPU引擎快469倍。该方法保持精确检索质量,在MS MARCO及三个BEIR基准测试中排名与参考版MaxSim一致。作为ColBERTv2/PLAID的即插即用替代方案,10万候选文档评分延迟从268毫秒降至1.2毫秒(降低98%端到端延迟)。进一步验证了从10万到50万文档的恒定吞吐、数据并行多卡划分、64-768维的鲁棒性,以及支持FP16/BF16/FP32。同时工作独立提出类似方案,本文差异在于针对维度大于128的分块设计和融合量化评分。
原文摘要 · Abstract (English)
Multi-vector retrieval models such as ColBERT achieve state-of-the-art accuracy through fine-grained token-level MaxSim scoring, yet existing GPU implementations leave most hardware performance unused. We give a roofline analysis of MaxSim on modern GPUs and identify a severe bandwidth gap: naive implementations reach only 5-18% of peak HBM bandwidth because they materialize the Nq x Nd similarity matrix, wasting memory traffic on data that is consumed once and discarded. We present TileMaxSim, a family of IO-aware Triton kernels that close this gap via (1) multi-query SRAM tiling that streams document embeddings through shared memory while accumulating per-query-token maxima in registers, reading each embedding from HBM exactly once; (2) dimension tiling that partitions the embedding dimension into 128-wide chunks, enabling scoring for d > 128 embeddings that overflow shared memory; and (3) fused product-quantization scoring via shared-memory lookup tables, cutting HBM I/O by up to ~31x. On NVIDIA H100 GPUs, TileMaxSim reaches 80.2% of peak HBM bandwidth and scores 82M documents/second (71.6M/s on real MS MARCO passages), a 220x speedup over loop-based scoring, 6.5x over fused PyTorch, 6.6-8.5x over torch.compile, and 469x the scoring throughput of WARP's CPU engine on the same node. TileMaxSim preserves exact retrieval quality: on MS MARCO and three BEIR benchmarks, rankings match reference MaxSim. As a drop-in replacement in ColBERTv2/PLAID, it cuts scoring latency at 100K candidates from 268 ms to 1.2 ms (98% lower end-to-end latency). We further show constant throughput from 100K to 500K documents, data-parallel multi-GPU sharding, robustness across dimensions 64-768, and FP16/BF16/FP32 support. Concurrent work independently develops an IO-aware fused MaxSim kernel; we differ in dimension tiling for d > 128 and fused product-quantization scoring.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。