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

Multi-vector retrieval models like ColBERT rely on the MaxSim operator—a fine-grained token-level similarity calculation—to achieve state-of-the-art search accuracy. However, existing GPU implementations are inefficient, often achieving only 5–18% of peak memory bandwidth because they materialize large, intermediate similarity matrices in HBM (High Bandwidth Memory). The author investigates whether IO-aware kernel design, which has successfully optimized attention mechanisms, can be applied to the unique matmul-max-sum reduction structure of MaxSim to eliminate this bottleneck.

The author introduces TileMaxSim, a suite of Triton kernels designed to minimize HBM traffic by keeping data in registers and shared memory (SRAM). The approach centers on three key optimizations:

• Multi-query SRAM tiling: Streams document embeddings through shared memory while accumulating per-query-token maxima in registers, ensuring each embedding is read from HBM exactly once.
• Dimension tiling: Partitions embedding dimensions into 128-wide chunks, allowing the kernel to handle arbitrary embedding sizes (up to 768) without exceeding register or SRAM capacity.
• Fused Product Quantization (PQ) scoring: Uses shared-memory lookup tables to score compressed documents directly, avoiding the need for explicit decompression and reducing HBM I/O by up to 31x.

TileMaxSim significantly outperforms existing baselines by aligning the computation with the GPU memory hierarchy. On NVIDIA H100 GPUs, it reaches 80.2% of peak HBM bandwidth, effectively transforming the MaxSim operation from a memory-bound bottleneck into a high-throughput process. In practical retrieval pipelines like ColBERTv2/PLAID, TileMaxSim acts as a drop-in replacement that reduces end-to-end scoring latency by 98% (from 268 ms to 1.2 ms for 100K candidates). Crucially, these performance gains are achieved while maintaining identical retrieval quality to reference implementations, as the kernel computes exact MaxSim scores.