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

arXiv: https://arxiv.org/abs/2606.26439 · HTML: https://mlsanity.com/paper/2606.26439/

## Key facts

- Cited by 1 papers in the map (data through 14 Oct 2026).
- Builds on 18 of its 30 references that are in the map.
- First posted on arXiv 24 Jun 2026.

## Builds on (most cited first, 18 of 18)

- [FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness](https://mlsanity.com/paper/2205.14135/) (2022)
- [FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioning](https://mlsanity.com/paper/2307.08691/) (2023)
- [Billion-scale similarity search with GPUs](https://mlsanity.com/paper/1702.08734/) (2017)
- [ColBERT: Efficient and Effective Passage Search via Contextualized Late Interaction over BERT](https://mlsanity.com/paper/2004.12832/) (2020)
- [ColBERTv2: Effective and Efficient Retrieval via Lightweight Late Interaction](https://mlsanity.com/paper/2112.01488/) (2021)
- [PLAID: An Efficient Engine for Late Interaction Retrieval](https://mlsanity.com/paper/2205.09707/) (2022)
- [FlashFFTConv: Efficient Convolutions for Long Sequences with Tensor Cores](https://mlsanity.com/paper/2311.05908/) (2023)
- [Rethinking the Role of Token Retrieval in Multi-Vector Retrieval](https://mlsanity.com/paper/2304.01982/) (2023)
- [AutoKernel: Autonomous GPU Kernel Optimization via Iterative Agent-Driven Search](https://mlsanity.com/paper/2603.21331/) (2026)
- [WARP: An Efficient Engine for Multi-Vector Retrieval](https://mlsanity.com/paper/2501.17788/) (2025)
- [FlashSinkhorn: IO-Aware Entropic Optimal Transport on GPU](https://mlsanity.com/paper/2602.03067/) (2026)
- [Sculpting the Vector Space: Towards Efficient Multi-Vector Visual Document Retrieval via Prune-then-Merge Framework](https://mlsanity.com/paper/2602.19549/) (2026)
- [ESPN: Memory-Efficient Multi-Vector Information Retrieval](https://mlsanity.com/paper/2312.05417/) (2023)
- [Scaling Attention via Feature Sparsity](https://mlsanity.com/paper/2603.22300/) (2026)
- [FLASH-MAXSIM: IO-Aware Fused Kernels for Late-Interaction Retrieval](https://mlsanity.com/paper/2605.29517/) (2026)
- [ColBERT-serve: Efficient Multi-Stage Memory-Mapped Scoring](https://mlsanity.com/paper/2504.14903/) (2025)
- [GEM: A Native Graph-based Index for Multi-Vector Retrieval](https://mlsanity.com/paper/2603.20336/) (2026)
- [Sparton: Fast and Memory-Efficient Triton Kernel for Learned Sparse Retrieval](https://mlsanity.com/paper/2603.25011/) (2026)

## Built on by (most cited first, 1 of 1)

- [MIDR: Enrichment-Augmented Indexing for Multimodal Document Retrieval](https://mlsanity.com/paper/2609.01316/) (2026)

## Abstract (arXiv)

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.