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

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.

• The paper introduces a novel IO-aware GPU kernel suite that streams document embeddings through shared memory to achieve optimal HBM bandwidth utilization.
• It employs multi-query and dimension tiling to efficiently support embedding dimensions up to 768, yielding up to 469x throughput improvements.
• Fused product quantization integration reduces IO by 31x, ensuring exact MaxSim scoring and enabling practical, large-scale neural search.

TileMaxSim: IO-Optimal GPU MaxSim Scoring via Dimension Tiling and Fused Product Quantization

Overview and Contribution

TileMaxSim directly addresses the suboptimal hardware utilization endemic to existing multi-vector retrieval accelerators, particularly those built around the MaxSim scoring operator as used in architectures such as ColBERT and its derivatives. MaxSim computes fine-grained token-to-token similarity matrices between query and document embeddings, followed by per-query-token maximum reduction and an outer sum. Despite recent advances in multi-vector retrieval models, prior GPU implementations squander the intensive high-bandwidth memory (HBM) resources of modern accelerators—often achieving under 20% of the available bandwidth—primarily due to naïve intermediate data materialization and inefficient memory traffic patterns.

The core contribution of TileMaxSim is a family of IO-aware GPU kernels implemented in Triton, designed explicitly to maximize effective utilization of GPU HBM bandwidth through three principal strategies:

1. Multi-query SRAM tiling: TileMaxSim streams document embeddings through shared memory (SRAM) while holding per-query-token maxima in registers. This design ensures that each document embedding is read from HBM exactly once, completely avoiding extraneous IO for temporaries.
2. Dimension tiling: For embedding dimensions exceeding SRAM/register capacity, TileMaxSim partitions the computation along the feature dimension, supporting arbitrary $d \in [64, 768]$ and maintaining high efficiency when $d > 128$, a regime unsupported by previous approaches.
3. Fused product quantization (PQ) decompression+scoring: TileMaxSim-PQ integrates PQ ADC table lookup and the MaxSim reduction without explicit decompression, reducing HBM IO by up to ~31x versus standard decompress-then-score pipelines.

The result is a kernel suite that achieves 80.2% of peak HBM bandwidth on an NVIDIA H100, 469x throughput improvement over WARP’s optimized CPU engine, and a 220x speedup over loop-based GPU scoring, whilst guaranteeing exact retrieval quality across all major benchmarks.

IO-Optimal Kernel Architecture

Roofline Analysis and Bottleneck Identification

Empirical and analytic roofline analysis demonstrates that MaxSim’s arithmetic intensity is far below the compute-to-bandwidth crossover point (e.g., 32 FLOP/byte vs. 591 on H100), establishing memory throughput as the limiting factor. Naïve implementations suffer from unnecessary traffic due to the full similarity matrix materialization, reading/writing tensors that are subsequently subject to immediate reduction.

TileMaxSim’s central design unifies matmul, per-row maximum, and outer-sum reductions into a single-pass kernel with carefully orchestrated tiling. In the optimal configuration, query tokens are processed together in register/SRAM-resident blocks, and each document embedding is loaded only once from HBM across all query tokens—a pattern that provably minimizes IO cost.

Dimension and Query-Document Tiling Strategies

Three kernel variants were developed:

• V1 (Per-query-token tiling): Assigns one block per query token but results in $N_q$ redundant reads of each document embedding.
• V2 (Per-document tiling): Limits redundant loads but still suboptimal for wide $N_d$ due to sequential reduction.
• V2-MQ (Multi-query tiling): Tiles over both query and document tokens (with tile size $BQ$ for query tokens), such that with $BQ = N_q$, each document embedding is read from HBM exactly once. This achieves the IO lower bound and approaches the roofline limit for bandwidth utilization.

For wide embedding dimensions, a further dimension-tiling variant partitions the embeddings into 128-wide blocks, enabling support for $d > 128$ without register/SRAM overflow.

Fused PQ Scoring

PQ is widely used to compress embeddings in production due to its aggressive reduction in memory/storage. TileMaxSim-PQ fuses ADC table construction into the uppermost memory cache (L2/Shared memory), and directly computes scores from compressed PQ codes without ever materializing decompressed vectors. This achieves theoretical and empirical IO reductions of up to 31x compared to decompress-then-score approaches, with practical throughput improvement of up to 2.9x in batch regimes relevant for real retrieval workloads.

Empirical Results

Throughput and Bandwidth Utilization

TileMaxSim V2-MQ yields 80.2% of the H100’s peak HBM bandwidth (2,687 GB/s of 3,350 GB/s). On synthetic data ($N_q=32$, $N_d=128$, $d=128$, $d > 128$0K) it achieves 82M docs/s, and on real tasks (MS MARCO) it realizes 71.6M docs/s (70% of peak). Variance in achieved bandwidth (44–86%) across parameters is accounted for by register pressure, tile misalignment, and atomic reduction overheads.

Comparative Benchmarking

Competitor	Speedup (TileMaxSim)	Notes
PyTorch Loop	220x	Gross CPU/GPU underutilization in looped baseline
torch.compile	6.6–8.5x	Compilers fail to discover deep fusion of matmul→max→sum
PLAID GPU	1.7–1.9x	Fast cuBLAS baseline; TileMaxSim avoids redundant writes
WARP (CPU)	469x	Both on same H100/CPU node; bandwidth advantage compounded

TileMaxSim delivers identical retrieval results to the reference implementation on both in-domain (MS MARCO) and cross-domain (BEIR benchmarks), with absolute ranking stability to within floating-point rounding. When integrated into full pipelines (e.g., ColBERTv2/PLAID), end-to-end latency at 100K candidates drops from 268ms to 1.2ms—a 98% reduction.

Scaling, Robustness, and Multi-GPU Support

TileMaxSim maintains nearly constant throughput from 100K to 500K candidates (83M docs/s), exhibits robust scaling with embedding dimension (dimension tiling supports up to 768), query/document token counts, and data type (FP16/BF16/FP32), and is trivially data-parallel: scoring is naturally sharded across multi-GPU clusters with no inter-device communication beyond result aggregation.

Theoretical and Practical Implications

TileMaxSim provides strong evidence for the efficacy of kernel-level IO-aware design in memory-bound operators. For the retrieval community, TileMaxSim demonstrates that brute-force, exact, large-scale MaxSim scoring is now feasible on a single H100 GPU at throughputs that previously required aggressive approximation or pruning on the CPU.

In practical terms, this capability redefines the tradeoff space for first-stage retrieval engines:

• Aggressive candidate pruning can be relaxed, potentially increasing recall.
• High-quality, cross-domain multi-vector models can be deployed with negligible scoring latency.
• PQ-compressed indices can be scored directly at scale with negligible IO overhead.

Theoretically, TileMaxSim’s IO complexity is provably optimal for GPU memory hierarchies under the MaxSim reduction structure. It also highlights opportunities for further system-level optimizations, including improved support for variable-length documents and fully asynchronous end-to-end GPU-to-SSD pipelines, as well as future meta-kernel search via compiler/LLM-agent code generation.

Further, this work establishes a reference for future IO-aware kernel design in other memory-bound reduction patterns (e.g., sparse attention, optimal transport, Chamfer similarity).

Conclusion

TileMaxSim decisively closes the GPU roofline gap for MaxSim scoring. By streaming document embeddings through SRAM and fusing the reduction path—including PQ decompression—the kernel family achieves near-optimal bandwidth utilization and consistent, exact retrieval results across all relevant benchmarks and retrieval paradigms. These results have immediate implications for the architecture of high-throughput, high-quality neural search systems, suggesting that the combination of IO-aware kernel design and high-bandwidth GPUs renders brute-force, exact, multi-vector retrieval practical at scales previously accessible only to aggressively approximate or single-vector methodologies.