TileMaxSim: IO-Aware GPU Kernels for MaxSim

• TileMaxSim is an IO-aware family of GPU kernels designed for exact token-level MaxSim scoring in multi-vector retrieval, optimizing memory I/O on modern GPUs.
• It fuses matrix multiplication, rowwise maximization, and product quantization lookup into on-chip computation to eliminate expensive similarity matrix materialization.
• TileMaxSim achieves up to 82M documents/sec on NVIDIA H100, reducing candidate scoring latency by over 200× while preserving exact retrieval quality.

Searching arXiv for TileMaxSim and related MaxSim kernel papers. {"query":"TileMaxSim IO-Aware GPU MaxSim Scoring with Dimension Tiling and Fused Product Quantization", "max_results": 5} {"query":"FLASH-MAXSIM IO-Aware Fused Kernels for Late-Interaction Scoring", "max_results": 5} TileMaxSim is an IO-aware family of GPU kernels for exact token-level MaxSim scoring in multi-vector retrieval, designed for settings such as ColBERT in which a query and a document are represented as sets of token embeddings and scored through late interaction rather than a single pooled vector (Sharma, 24 Jun 2026). The central objective is to close a severe memory-bandwidth gap on modern GPUs by avoiding materialization of the $N_q \times N_d$ similarity matrix and instead fusing matrix multiplication, rowwise maximization, and score accumulation into SRAM- and register-resident kernels. In the reported implementation on an NVIDIA H100, TileMaxSim V2-MQ reaches $80\%$ of peak HBM3 bandwidth, scores up to $82$ M documents/sec ($71.6$ M/sec on real MS MARCO), preserves exact retrieval quality, and reduces ColBERTv2/PLAID’s $100$ K candidate scoring from $268$ ms to $1.2$ ms (Sharma, 24 Jun 2026).

1. Formal setting and MaxSim definition

TileMaxSim addresses token-level MaxSim scoring for multi-vector retrieval. If a query $q$ produces $N_q$ token embeddings $Q \in \mathbb{R}^{N_q \times d}$ and a document $80\%$0 produces $80\%$1 embeddings $80\%$2, the score is defined as

$80\%$3

with the equivalent formulation

$80\%$4

and

$80\%$5

For a batch of $80\%$6 documents, each dot product of dimension $80\%$7 costs $80\%$8 FLOPs plus one comparison for max, yielding

$80\%$9

The paper distinguishes sharply between arithmetic cost and IO cost. A naïve implementation that materializes $82$0 reads $82$1 and $82$2, and then writes and rereads the similarity matrix in FP32, giving

$82$3

A fused SRAM-tiled kernel avoids writing $82$4:

$82$5

For typical ColBERT parameters $82$6, the reported IO is approximately $82$7 GB for the naïve method and approximately $82$8 GB for TileMaxSim, with arithmetic intensity increasing from $82$9 to $71.6$0 FLOP/byte (Sharma, 24 Jun 2026). This establishes the basic computational object that TileMaxSim optimizes: not the algebraic definition of MaxSim itself, but the way its intermediate state is moved through the GPU memory hierarchy.

2. Roofline analysis and the bandwidth bottleneck

The motivating observation is that existing GPU implementations of MaxSim leave most hardware performance unused because they waste HBM traffic on similarities that are consumed once and discarded. On the H100, the paper reports $71.6$1 TFLOP/s FP16 peak and $71.6$2 TB/s HBM3 bandwidth, so the roofline crossover is $71.6$3 FLOP/byte; MaxSim’s arithmetic intensity of $71.6$4–$71.6$5 FLOP/byte lies far below that threshold, confirming that the workload is memory-bound (Sharma, 24 Jun 2026).

Within this regime, materializing the similarity matrix is especially costly. The paper states that naïve implementations achieve only $71.6$6–$71.6$7 of peak bandwidth, approximately $71.6$8 TB/s, because the $71.6$9 matrix is written to HBM and then reduced immediately afterward. TileMaxSim targets a $100$0–$100$1 IO reduction by eliminating that intermediate and restructuring the computation around on-chip storage.

The reported bandwidth measurements make the design objective concrete. On H100, achieved bandwidth utilization for PyTorch Naive is $100$2 at $100$3 K, $100$4 K, and $100$5 K, while PyTorch Loop reaches only $100$6–$100$7. By contrast, TileMaxSim V2-MQ reaches $100$8 at $100$9 K, $268$0 at $268$1 K, and $268$2 at $268$3 K and $268$4 K. At $268$5 K, V2-MQ reaches $268$6 TB/s, i.e. $268$7 of $268$8 TB/s, when $268$9 and each document embedding is read once (Sharma, 24 Jun 2026). A plausible implication is that the paper’s main contribution is best understood as an IO-optimal scheduling result expressed through Triton kernels rather than as a new retrieval scoring function.

3. Kernel architecture: multi-query tiling, dimension tiling, and fused PQ

The TileMaxSim kernel family is organized around three techniques: multi-query SRAM tiling, embedding-dimension tiling, and fused product-quantization lookup-table scoring (Sharma, 24 Jun 2026).

Component	Mechanism	Reported role
Multi-query SRAM tiling	Streams document tiles through shared memory and keeps per-query-token maxima in registers	Reads each embedding from HBM exactly once when $1.2$0
Embedding-dimension tiling	Partitions $1.2$1 into $1.2$2-wide chunks	Enables scoring for $1.2$3 embeddings that overflow shared memory
Fused PQ scoring	Uses shared-memory lookup tables instead of explicit decompression	Cuts HBM I/O by up to $1.2$4

Multi-query SRAM tiling

In TileMaxSim V2-MQ, each GPU program block is parameterized by $1.2$5, identifying a document batch index and a query-tile index of size $1.2$6. The kernel loads a query tile into registers, initializes a register-resident maximum array to $1.2$7, iterates over document tiles in shared memory, computes $1.2$8, updates rowwise maxima in registers, and finishes with an atomic_add of the summed maxima to the document score. The key point is that similarity never spills to HBM; all partial results remain in registers or SRAM.

The corresponding IO Complexity Theorem states that with $1.2$9 and tile sizes $q$0,

$q$1

bytes, so each embedding is read exactly once (Sharma, 24 Jun 2026). This is the paper’s strongest optimality claim.

Embedding-dimension tiling

When $q$2, full-width dot products exceed on-chip register/SRAM capacity. TileMaxSim therefore partitions the embedding dimension into $q$3 tiles of width $q$4 except possibly the last. For each dimension tile $q$5, the kernel loads slices of $q$6 and $q$7, computes partial similarities, and updates running maxima. The paper formalizes this through vectors $q$8 and $q$9, maintaining

$N_q$0

This scheme supports $N_q$1 up to $N_q$2 in the same kernel family and achieves $N_q$3 of peak bandwidth at $N_q$4 (Sharma, 24 Jun 2026). That result distinguishes TileMaxSim from implementations tuned only for the standard $N_q$5 ColBERT configuration.

Fused product-quantization scoring

For PQ, each $N_q$6-dimensional vector is split into $N_q$7 subspaces of size $N_q$8, and each subspace is quantized to one of $N_q$9 centroids $Q \in \mathbb{R}^{N_q \times d}$0. A document token is stored as codes $Q \in \mathbb{R}^{N_q \times d}$1, and asymmetric lookup-table scoring precomputes

$Q \in \mathbb{R}^{N_q \times d}$2

Then, for each document token $Q \in \mathbb{R}^{N_q \times d}$3,

$Q \in \mathbb{R}^{N_q \times d}$4

TileMaxSim-PQ fuses this into the MaxSim kernel. Phase 1 builds $Q \in \mathbb{R}^{N_q \times d}$5 in parallel across $Q \in \mathbb{R}^{N_q \times d}$6 programs; for $Q \in \mathbb{R}^{N_q \times d}$7, $Q \in \mathbb{R}^{N_q \times d}$8, and $Q \in \mathbb{R}^{N_q \times d}$9, the lookup table occupies approximately $80\%$00 KB and fits in L2 cache. Phase 2 loads PQ codes for each document tile and performs $80\%$01 lookups per token from $80\%$02 in shared memory or L2, without HBM writes of decompressed vectors. For $80\%$03 K, $80\%$04, $80\%$05, $80\%$06, $80\%$07, HBM IO decreases from approximately $80\%$08 GB for decompress-and-score to approximately $80\%$09 GB for TileMaxSim-PQ, a $80\%$10 reduction (Sharma, 24 Jun 2026).

4. Performance characteristics and baseline comparisons

The paper reports measurements via CUDA events over $80\%$11 runs on a single H100 for $80\%$12, $80\%$13, $80\%$14, and $80\%$15 up to $80\%$16 K (Sharma, 24 Jun 2026). Under these conditions, the principal throughput results are as follows.

Configuration	Baseline	TileMaxSim result
$80\%$17, $80\%$18 K	PyTorch Loop: $80\%$19 M/s	V2-MQ: $80\%$20 M/s
$80\%$21, $80\%$22 K	PyTorch Naive: $80\%$23 M/s	V2-MQ: $80\%$24 M/s
$80\%$25, $80\%$26 K	PyTorch Loop: $80\%$27 M/s	V2-MQ: $80\%$28 M/s
$80\%$29, $80\%$30 K	PyTorch Naive: $80\%$31 M/s	V2-MQ: $80\%$32 M/s

These correspond to speedups of $80\%$33 over loop-based scoring at $80\%$34, $80\%$35 K; $80\%$36 over loops and $80\%$37 over PyTorch Naive at $80\%$38, $80\%$39 K; $80\%$40 over loops at $80\%$41, $80\%$42 K; and $80\%$43 over loops and $80\%$44 over PyTorch Naive at $80\%$45, $80\%$46 K (Sharma, 24 Jun 2026).

Against GPU baselines that are closer to deployed retrieval codepaths, TileMaxSim remains faster. Versus ColBERTv2/PLAID’s colbert_score, described as batched GEMM $80\%$47 materialize $80\%$48 $80\%$49 max $80\%$50 sum, it is $80\%$51–$80\%$52 faster for $80\%$53 K. Versus torch.compile(mode="max-autotune"), it is $80\%$54–$80\%$55 faster because the compiler fails to fuse matmul $80\%$56 max $80\%$57 sum. For $80\%$58, $80\%$59, and $80\%$60 K, the reported latencies are $80\%$61 ms for PLAID GPU, $80\%$62 ms for torch.compile, and $80\%$63 ms for TileMaxSim V2-MQ (Sharma, 24 Jun 2026).

The PQ path also shows measurable throughput gains. For $80\%$64, $80\%$65 K, decompress-and-score reaches $80\%$66 M/s and TileMaxSim-PQ reaches $80\%$67 M/s; for $80\%$68, $80\%$69 K, the figures are $80\%$70 M/s and $80\%$71 M/s; for $80\%$72, $80\%$73 K, they are $80\%$74 M/s and $80\%$75 M/s. The corresponding speedups are $80\%$76, $80\%$77, and $80\%$78 (Sharma, 24 Jun 2026).

The CPU comparison is included primarily as a system-level reference point. The WARP CPU engine on Xeon DDR5, single-threaded, scores approximately $80\%$79 K docs/s in $80\%$80 ms for a bound of at most $80\%$81 docs/query, whereas TileMaxSim scores $80\%$82 K docs in $80\%$83 ms, i.e. $80\%$84 M/s, a $80\%$85 throughput gap (Sharma, 24 Jun 2026). The paper attributes this to the $80\%$86 memory-bandwidth difference and the GPU’s $80\%$87 bandwidth utilization.

5. Integration into retrieval pipelines and scalability

TileMaxSim is presented as a drop-in replacement for the scoring kernel in ColBERTv2/PLAID, with no changes required in indexing or candidate generation and with identical retrieval quality (Sharma, 24 Jun 2026). In the PLAID integration experiment, scoring $80\%$88 K candidates decreases from $80\%$89 ms to $80\%$90 ms, a $80\%$91 speedup, and scoring $80\%$92 K candidates decreases from $80\%$93 ms to $80\%$94 ms, a $80\%$95 speedup. The paper summarizes the latter as a $80\%$96 end-to-end latency reduction.

Quality preservation is a central part of the integration claim. On MS MARCO dev and three BEIR datasets, the paper reports identical rankings, with maximum $80\%$97 and identical MRR@$80\%$98 and nDCG@$80\%$99 (Sharma, 24 Jun 2026). A common misconception is that aggressive kernel fusion or mixed-precision support necessarily changes retrieval behavior; in the reported experiments, exact MaxSim rankings are preserved.

Scalability is described along several axes. Throughput saturates at approximately $82$00 M docs/s for $82$01 K and remains constant up to $82$02 K, specifically $82$03–$82$04 M/s. Multi-GPU sharding is characterized as embarrassingly parallel across documents, with near-linear scaling across multiple H100s using simple data-parallel splits. Embedding dimensionality from $82$05 to $82$06 is supported through dimension tiling, with reported throughputs of $82$07 M/s at $82$08 ($82$09 bandwidth), $82$10 M/s at $82$11 ($82$12 bandwidth), and $82$13 M/s at $82$14 ($82$15 bandwidth). The implementation supports FP16, BF16, and FP32; FP16 and BF16 yield identical approximately $82$16 M/s at $82$17, $82$18 K, whereas FP32 is $82$19 slower due to doubled IO (Sharma, 24 Jun 2026).

These results indicate that the kernel family is intended not merely as a microbenchmark artifact but as an operational scoring primitive for large candidate sets. This suggests that the main system value lies in making brute-force multi-vector scoring practical at scales where the scoring stage previously dominated latency.

6. Position within IO-aware late-interaction scoring research

TileMaxSim belongs to a broader line of IO-aware fused MaxSim kernels for late-interaction retrieval. The paper explicitly notes concurrent work that independently develops an IO-aware fused MaxSim kernel, and states that TileMaxSim differs in dimension tiling for $82$20 and fused product-quantization scoring (Sharma, 24 Jun 2026). That concurrent work is "FLASH-MAXSIM: IO-Aware Fused Kernels for Late-Interaction Scoring" (Pony et al., 28 May 2026).

The relationship between the two works is technically specific. FLASH-MAXSIM also avoids materializing the full query-token $82$21 document-token similarity tensor, instead streaming tiles through on-chip SRAM and folding the row-maximum reduction into the same pass. It extends the IO-aware principle to the training backward pass via an inverse-grid CSR construction, and further includes INT8$82$22INT8 quantization and padding-free scoring (Pony et al., 28 May 2026). TileMaxSim, by contrast, emphasizes a family of Triton kernels targeted at exact inference-time MaxSim scoring, with explicit support for embedding-dimension tiling up to $82$23 and a fused PQ path (Sharma, 24 Jun 2026).

This proximity of contributions clarifies both the maturity and the scope of the area. The shared premise is that late-interaction retrieval is constrained less by raw Tensor Core throughput than by avoidable HBM traffic; the principal differences are in which parts of the retrieval or training stack are fused and which deployment constraints are prioritized. A plausible implication is that IO-aware MaxSim has become a distinct optimization layer within multi-vector retrieval systems, analogous to the role of FlashAttention in dense sequence modeling, although the paper itself does not use that analogy.