---
title: 'TileMaxSim: IO-Aware GPU Kernels for MaxSim'
url: https://www.emergentmind.com/topics/tilemaxsim
type: topic
---

# TileMaxSim: IO-Aware GPU Kernels for MaxSim

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 [2606.26439]. 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 [2606.26439].

## 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 $p$ produces $N_d$ embeddings $D \in \mathbb{R}^{N_d \times d}$, the score is defined as

$$
\mathrm{MaxSim}(Q,D) = \sum_{i=1}^{N_q} \max_{j=1}^{N_d} q_i \cdot d_j
$$

with the equivalent formulation

$$
S = QD^\top \in \mathbb{R}^{N_q \times N_d}
$$

and

$$
\mathrm{score} = \sum_{i=1}^{N_q} \max_{j=1}^{N_d} S_{i,j}.
$$

For a batch of $B$ documents, each dot product of dimension $d$ costs $2d$ FLOPs plus one comparison for max, yielding

$$
\mathrm{FLOPs} = B \cdot N_q \cdot N_d \cdot (2d + 1).
$$

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

$$
\mathrm{IO}_{\text{naive}} = 2 \cdot N_q \cdot d + 2 \cdot B \cdot N_d \cdot d + 8 \cdot B \cdot N_q \cdot N_d.
$$

A fused SRAM-tiled kernel avoids writing $S$:

$$
\mathrm{IO}_{\text{fused}} = 2 \cdot N_q \cdot d + 2 \cdot B \cdot N_d \cdot d + 4 \cdot B \cdot N_q.
$$

For typical ColBERT parameters $(N_q = 32, N_d = 128, d = 128, B = 10{,}000)$, the reported IO is approximately $0.655$ GB for the naïve method and approximately $0.329$ GB for TileMaxSim, with arithmetic intensity increasing from $16.1$ to $32.0$ FLOP/byte [2606.26439]. 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 $1979$ TFLOP/s FP16 peak and $3.35$ TB/s HBM3 bandwidth, so the roofline crossover is $591$ FLOP/byte; MaxSim’s arithmetic intensity of $16$–$32$ FLOP/byte lies far below that threshold, confirming that the workload is memory-bound [2606.26439].

Within this regime, materializing the similarity matrix is especially costly. The paper states that naïve implementations achieve only $5$–$18\%$ of peak bandwidth, approximately $0.17$ TB/s, because the $N_q \times N_d$ matrix is written to HBM and then reduced immediately afterward. TileMaxSim targets a $2$–$3\times$ 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 $18\%$ at $B=1$ K, $10$ K, and $50$ K, while PyTorch Loop reaches only $0.3$–$0.4\%$. By contrast, TileMaxSim V2-MQ reaches $55.9\%$ at $B=1$ K, $78.0\%$ at $B=10$ K, and $80.2\%$ at $B=50$ K and $100$ K. At $B=100$ K, V2-MQ reaches $2.69$ TB/s, i.e. $80\%$ of $3.35$ TB/s, when $B_Q = N_q$ and each document embedding is read once [2606.26439]. 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 [2606.26439].

| 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 $B_Q=N_q$ |
| Embedding-dimension tiling | Partitions $d$ into $128$-wide chunks | Enables scoring for $d>128$ embeddings that overflow shared memory |
| Fused PQ scoring | Uses shared-memory lookup tables instead of explicit decompression | Cuts HBM I/O by up to $\sim 31\times$ |

### Multi-query SRAM tiling

In TileMaxSim V2-MQ, each GPU program block is parameterized by $(b, q_b)$, identifying a document batch index and a query-tile index of size $B_Q$. The kernel loads a query tile into registers, initializes a register-resident maximum array to $-\infty$, iterates over document tiles in shared memory, computes $S_{\text{tile}} = \mathrm{tl.dot}(Q_{\text{tile}}, D_{\text{tile}}^\top)$, 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 $B_Q=N_q$ and tile sizes $B_Q, B_N$,

$$
\mathrm{IO}_{\text{optimal}} = 2 \cdot (N_q \cdot d + B \cdot N_d \cdot d) + 4 \cdot B
$$

bytes, so each embedding is read exactly once [2606.26439]. This is the paper’s strongest optimality claim.

### Embedding-dimension tiling

When $d > 128$, full-width dot products exceed on-chip register/SRAM capacity. TileMaxSim therefore partitions the embedding dimension into $T$ tiles of width $D_t = 128$ except possibly the last. For each dimension tile $\ell$, the kernel loads slices of $Q_{\text{tile}}$ and $D_{\text{tile}}$, computes partial similarities, and updates running maxima. The paper formalizes this through vectors $q_i^{(\ell)} \in \mathbb{R}^{D_t}$ and $d_j^{(\ell)} \in \mathbb{R}^{D_t}$, maintaining

$$
m_i = \max_j \sum_{\ell} \left(q_i^{(\ell)} \cdot d_j^{(\ell)}\right).
$$

This scheme supports $d$ up to $768$ in the same kernel family and achieves $86\%$ of peak bandwidth at $d=768$ [2606.26439]. That result distinguishes TileMaxSim from implementations tuned only for the standard $d=128$ ColBERT configuration.

### Fused product-quantization scoring

For PQ, each $d$-dimensional vector is split into $M$ subspaces of size $d_{\text{sub}} = d/M$, and each subspace is quantized to one of $K$ centroids $C[m,k] \in \mathbb{R}^{d_{\text{sub}}}$. A document token is stored as codes $c_{j,1..M}$, and asymmetric lookup-table scoring precomputes

$$
T[i,m,k] = q_i[m \cdot d_{\text{sub}} : (m+1) \cdot d_{\text{sub}}] \cdot C[m,k].
$$

Then, for each document token $j$,

$$
\mathrm{score}_{i,j} \approx \sum_{m=1}^{M} T[i,m,c_{j,m}].
$$

TileMaxSim-PQ fuses this into the MaxSim kernel. Phase 1 builds $T$ in parallel across $N_q \times M$ programs; for $N_q=32$, $M=16$, and $K=256$, the lookup table occupies approximately $512$ KB and fits in L2 cache. Phase 2 loads PQ codes for each document tile and performs $M$ lookups per token from $T$ in shared memory or L2, without HBM writes of decompressed vectors. For $B=100$ K, $N_q=32$, $N_d=128$, $M=16$, $K=256$, HBM IO decreases from approximately $6.76$ GB for decompress-and-score to approximately $0.22$ GB for TileMaxSim-PQ, a $31\times$ reduction [2606.26439].

## 4. Performance characteristics and baseline comparisons

The paper reports measurements via CUDA events over $50$ runs on a single H100 for $N_q=32$, $d=128$, $N_d \in \{64,128,256\}$, and $B$ up to $100$ K [2606.26439]. Under these conditions, the principal throughput results are as follows.

| Configuration | Baseline | TileMaxSim result |
|---|---|---|
| $N_d=128$, $B=100$ K | PyTorch Loop: $0.37$ M/s | V2-MQ: $81.2$ M/s |
| $N_d=128$, $B=10$ K | PyTorch Naive: $9.4$ M/s | V2-MQ: $60.4$ M/s |
| $N_d=64$, $B=100$ K | PyTorch Loop: $0.87$ M/s | V2-MQ: $121.8$ M/s |
| $N_d=256$, $B=10$ K | PyTorch Naive: $5.1$ M/s | V2-MQ: $40.0$ M/s |

These correspond to speedups of $220\times$ over loop-based scoring at $N_d=128$, $B=100$ K; $74\times$ over loops and $6.4\times$ over PyTorch Naive at $N_d=128$, $B=10$ K; $140\times$ over loops at $N_d=64$, $B=100$ K; and $182\times$ over loops and $7.8\times$ over PyTorch Naive at $N_d=256$, $B=10$ K [2606.26439].

Against GPU baselines that are closer to deployed retrieval codepaths, TileMaxSim remains faster. Versus ColBERTv2/PLAID’s `colbert_score`, described as batched GEMM $\rightarrow$ materialize $S$ $\rightarrow$ max $\rightarrow$ sum, it is $1.7$–$1.9\times$ faster for $B \ge 10$ K. Versus `torch.compile(mode="max-autotune")`, it is $6.6$–$8.5\times$ faster because the compiler fails to fuse matmul $\rightarrow$ max $\rightarrow$ sum. For $N_q=32$, $N_d=128$, and $B=100$ K, the reported latencies are $2.06$ ms for PLAID GPU, $4.44$ ms for `torch.compile`, and $1.19$ ms for TileMaxSim V2-MQ [2606.26439].

The PQ path also shows measurable throughput gains. For $N_d=128$, $B=10$ K, decompress-and-score reaches $2.9$ M/s and TileMaxSim-PQ reaches $5.6$ M/s; for $N_d=256$, $B=10$ K, the figures are $1.5$ M/s and $2.9$ M/s; for $N_d=64$, $B=1$ K, they are $1.5$ M/s and $4.4$ M/s. The corresponding speedups are $1.9\times$, $1.9\times$, and $2.9\times$ [2606.26439].

The CPU comparison is included primarily as a system-level reference point. The WARP CPU engine on Xeon DDR5, single-threaded, scores approximately $175$ K docs/s in $1.12$ ms for a bound of at most $196$ docs/query, whereas TileMaxSim scores $500$ K docs in $1.23$ ms, i.e. $82$ M/s, a $469\times$ throughput gap [2606.26439]. The paper attributes this to the $43\times$ memory-bandwidth difference and the GPU’s $80\%$ 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 [2606.26439]. In the PLAID integration experiment, scoring $10$ K candidates decreases from $23.5$ ms to $0.17$ ms, a $138\times$ speedup, and scoring $100$ K candidates decreases from $267.4$ ms to $1.21$ ms, a $220\times$ speedup. The paper summarizes the latter as a $98\%$ 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 $\Delta \mathrm{score} < 10^{-5}$ and identical MRR@$n$ and nDCG@$n$ [2606.26439]. 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 $83$ M docs/s for $B \ge 100$ K and remains constant up to $500$ K, specifically $81$–$83$ 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 $64$ to $768$ is supported through dimension tiling, with reported throughputs of $90.5$ M/s at $d=64$ ($44\%$ bandwidth), $60.9$ M/s at $d=128$ ($60\%$ bandwidth), and $14.7$ M/s at $d=768$ ($86\%$ bandwidth). The implementation supports FP16, BF16, and FP32; FP16 and BF16 yield identical approximately $60$ M/s at $d=128$, $B=10$ K, whereas FP32 is $1.6\times$ slower due to doubled IO [2606.26439].

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 $d>128$ and fused product-quantization scoring [2606.26439]. That concurrent work is "FLASH-MAXSIM: IO-Aware Fused Kernels for Late-Interaction Scoring" [2605.29517].

The relationship between the two works is technically specific. FLASH-MAXSIM also avoids materializing the full query-token $\times$ 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$\times$INT8 quantization and padding-free scoring [2605.29517]. 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 $d=768$ and a fused PQ path [2606.26439].

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.

Source: https://www.emergentmind.com/topics/tilemaxsim