TileMaxSim: IO-Aware GPU MaxSim Scoring with Dimension Tiling and Fused Product Quantization

2026-06-24

AI Summary

This work addresses the severe memory bandwidth underutilization (only 5–18%) in existing GPU implementations of multi-vector retrieval models like ColBERT during MaxSim scoring, which critically limits efficiency. The authors propose an I/O-aware Triton kernel that integrates multi-query SRAM tiling, 128-dimensional chunking, and fused product quantization to reduce HBM I/O by 31× while preserving exact MaxSim results and ranking quality. Evaluated on an NVIDIA H100 GPU, the method achieves 80.2% of peak memory bandwidth utilization, enabling a throughput of 82 million documents per second—6.5× faster than a PyTorch baseline—and reduces end-to-end latency by 98%. Full ranking fidelity is maintained, with identical retrieval performance to the original model on both MS MARCO and BEIR benchmarks.

Abstract

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.

Problem

• MaxSim
• GPU bandwidth
• multi-vector retrieval
• memory efficiency
• similarity scoring

Innovation

• IO-aware
• dimension tiling
• fused product quantization
• MaxSim
• GPU kernel