1 Introduction↩︎

Dense single-vector retrieval encodes each query and document as one vector and scores by a dot product. Late-interaction retrieval instead keeps a set of token-level embeddings per item and scores with MaxSim: \[\label{eq:maxsim} \mathrm{score}(Q, D) = \sum_{i=1}^{L_q} \max_{j=1}^{L_d} \langle Q_i, D_j \rangle,\tag{1}\] with \(Q \in \mathbb{R}^{L_q \times d}\), \(D \in \mathbb{R}^{L_d \times d}\). This finer-grained matching is what gives ColBERT [1] and its document-image successor ColPali [2] their quality advantage, and its cost. The natural PyTorch implementation forms the \([L_q, L_d]\) similarity matrix per (query, document) pair via an einsum, then reduces it (\(\max\) over document tokens, sum over query tokens). At scale the dominant memory cost is therefore a \([B, L_q, L_d]\) tensor that exists only to be reduced away.

For visual ColPali (\(L_q\!=\!L_d\!=\!1024\) patch tokens, \(d\!=\!128\)) over a \(10\mathrm{K}\)-document candidate set, that tensor is \(10^4 \cdot 1024^2 \cdot 2 \approx 21\) GB in FP16 (\(42\) GB in FP32). It out-of-memories (OOMs) a \(40\) GB GPU and consumes a quarter to half of an \(80\) GB one, capping batch size for reranking, for in-batch-negatives contrastive training, and for distillation. The matrix is the bottleneck, and it never needed to exist.

FlashAttention [3] showed that the analogous attention matrix can be avoided by IO-aware tiling. We make the same move for MaxSim; and because MaxSim is an operator, not just an inference pass, we carry the idea through training and quantization. Our contributions:

1. A fused forward kernel (§4.1) that computes exact MaxSim with peak memory \(\approx\) the document embeddings, the similarity tensor never forms. Up to \(3.9\times\) over naive PyTorch on dense scoring (A100; \(4.7\times\) on H100) and OOM-unlock at corpus scale (§5).

2. A fused backward kernel (§4.2): an inverse-grid CSR map, built at runtime from the saved forward argmax, that performs an atomic-free, destination-owned gradient reduction, no re-materialized matrix, FP32-accurate gradients. It unlocks contrastive batches that OOM PyTorch.

3. Quantized and variable-length variants (§4.3): INT8\(\times\)INT8 MaxSim with Spearman \(\rho\!=\!0.999\) ranking fidelity, and padding-free cu_seqlens scoring (up to \(5\times\)) that removes the padding tax.

4. An open-source release of all kernels and variants.

2 Related Work↩︎

FlashAttention [3] is the conceptual parent, IO-aware tiling that avoids materializing the attention matrix by streaming tiles through on-chip memory with an online softmax. Flash-MaxSim applies the same principle to the MaxSim operator; the reduction is a plain running max (idempotent, no rescaling), and, unlike attention kernels, we carry it through a hard-selection (argmax) backward for training. Flash-KMeans [4] is the closest analog in spirit: a materialization-free assignment via online argmin plus a sort-inverse, low-contention centroid update; we target a retrieval operator and add a quantization and variable-length story. ColBERT [1] and ColPali [2] define the late-interaction application this work accelerates.

3 Background and Motivation↩︎

3.1 MaxSim and the Standard GPU Implementation↩︎

A late-interaction model encodes a query as \(Q \in \mathbb{R}^{L_q\times d}\) and a document as \(D \in \mathbb{R}^{L_d\times d}\) (token-level, \(\ell_2\)-normalized, \(d{=}128\)). MaxSim (Eq. 1 ) is computed at scale in two regimes: reranking (one query vs.\(B\) candidate documents, output \([B]\)) and, during training, in-batch-negatives scoring (\(N_q\) queries vs. \(B\) documents, an all-pairs \([N_q, B]\) score matrix). Written directly,

S = einsum("bqd,bld->bql", Q, D)
#   [B, Lq, Ld]: materialized in HBM
score = S.max(dim=2).values.sum(dim=1)

the intermediate \(S\) of shape \([B, L_q, L_d]\) is written to high-bandwidth memory (HBM), read back for the \(\max\) reduction, and discarded.

3.2 Kernel-Level Bottleneck: the Materialized Tensor↩︎

The footprint of \(S\) is \(O(B \cdot L_q \cdot L_d)\), independent of \(d\), and it is both the memory and the bandwidth bottleneck. The naive path moves \(\Theta(B L_q L_d)\) bytes just for \(S\) (one write, at least one read), versus \(\Theta(B L_d d)\) bytes for the inputs. For ColPali (\(L_q{=}L_d{=}1024\), \(d{=}128\)) the \(S\) traffic exceeds the input traffic by \(L_q/d \approx 8\times\), and the materialized \(S\) exceeds the inputs by \(L_q \approx 1024\times\).

3.2.0.1 Arithmetic intensity and the roofline.

Scoring one pair costs \(2 L_q L_d d\) FLOPs. The naive path’s HBM traffic is dominated by the \(L_q\times L_d\) surface of \(S\) (\(\approx 4 L_q L_d\) bytes in FP16: one write, two reads), so its arithmetic intensity is \(\mathrm{AI}_{\text{naive}} \approx 2 L_q L_d d / (4 L_q L_d) = d/2 \approx 64\) FLOPs/byte at \(d{=}128\). An NVIDIA H100’s bf16 roofline ridge point is \(\approx 295\) FLOPs/byte, so the naive operator sits four-to-five times into the memory-bound regime, the tensor cores idle while HBM moves a tensor that is immediately discarded. Flash-MaxSim removes the \(S\) traffic entirely and crosses the ridge (§4.1).

3.3 System-Level Constraints in Real Deployments↩︎

The materialized tensor imposes hard system limits, not just slowdowns. At ColPali corpus scale the \(21\) GB FP16 \(S\) OOMs a \(40\) GB GPU and consumes half of an \(80\) GB one, so reranking a large candidate set is impossible without manual chunking. In training, in-batch-negatives scoring forms an all-pairs \([N_q, B, L_q, L_d]\) tensor whose size is quadratic in the batch \(B\); with autograd also retaining its gradient, the contrastive batch size is capped far below what the model and optimizer alone would allow, at ColPali shape, \(B{=}128\) OOMs an \(80\) GB GPU (§5). These caps, on corpus size at inference and batch size at training, are the real-deployment cost of materialization.

4 Methodology↩︎

4.1 Fused Forward: Materialization-Free Scoring via Online Max↩︎

The forward kernel computes Eq. 1 without ever writing \(S\) (Figure 2). It is structurally a FlashAttention-style tile-then-reduce, with three differences that make it cheaper: (i) the in-tile reduction is a plain running maximum, not a softmax, no log-sum-exp and no rescaling, because \(\max\) is idempotent; (ii) the output is a scalar per (query, document) pair, not a matrix; (iii) there is no value matrix \(V\) to re-read.

The kernel launches \(N_q B\) independent programs (thread blocks), one per (query, document) pair; each produces a single FP32 scalar. The similarity sub-tile \(S_t\) of shape \([\mathrm{\small b}_q, \mathrm{\small b}_d]\) lives only in SRAM (the running max \(m\) and accumulator stay in fp32 registers) and is reduced to a \(\mathrm{\small b}_q\)-vector before the next document tile overwrites it. The full \([B, L_q, L_d]\) tensor is never written. Because the sum-of-maxima decomposes over query chunks, a long \(L_q\) is split into fixed chunks whose partial scores are summed, so a single compiled kernel serves all problem shapes at runtime.

The running-max reduction is idempotent and associative, so partial tiles compose by a single elementwise \(\max\) with no rescaling correction, unlike the online softmax of FlashAttention (detailed next).

4.1.1 Online-Max Recurrence↩︎

After consuming the \(k\)-th document tile, the running max obeys \(m^{(k)} = \max(m^{(k-1)}, \mathrm{rowmax}(S^{(k)}))\) with \(m^{(0)} = -\infty\); after all tiles, \(m\) is the exact per-row maximum over the full \(L_d\) axis. Unlike online softmax (FlashAttention), there is no running normalizer and no rescaling factor \(\exp(m^{(k-1)}{-}m^{(k)})\): \(\max\) is idempotent, so the update is a single elementwise max and the result is bitwise the offline computation. Masked document positions are set to \(-\infty\) before the row reduction (not post-multiplied by a \(0/1\) mask), so padding can never be selected even when valid scores are negative.

4.1.2 IO Complexity and Roofline↩︎

Proof sketch. Each program loads its \(Q\) and \(D_b\) tiles once into SRAM, forms \(S_t\) on chip, and folds it into the running max; \(S\) is never written to or read from HBM, so operand traffic is \(\Theta(N_q B (L_q{+}L_d)d)\) and output traffic \(\Theta(N_q B)\). The naive path additionally writes and reads \(S\), adding \(\Theta(N_q B L_q L_d)\), which dominates whenever \(L_q L_d \gg (L_q{+}L_d)d\). \(\square\)

With only the operand reads as HBM traffic, \(\mathrm{AI}_{\text{fused}} \approx 2 L_q L_d d / (2(L_q{+}L_d)d) = L_q L_d/(L_q{+}L_d) \approx 512\) FLOPs/byte at ColPali shape, above the H100 ridge. The identical matmul thus moves from memory-bound to compute-bound: it is faster not because the arithmetic is cheaper, but because it no longer waits on memory.

4.1.3 Numerical Fidelity↩︎

The accumulator is FP32 throughout. The maximum relative error against an FP32 reference is \(2\times10^{-6}\), versus \(4\times10^{-3}\) for an FP16 batched-matrix-multiply baseline, three orders of magnitude tighter, so ranking ties are resolved at true precision rather than by FP16 noise.

4.1.4 Kernel Family↩︎

A family of forward variants shares the running-max core: single-query rerank, batched multi-query, variable-length packed (§4.3), query-reuse, split-\(K\), and a two-stage INT8\(\to\)FP16 top-\(K\) scan, selected by a runtime dispatcher on \((N_q, B, L_q, L_d, d, \mathrm{dtype})\).

4.2 Inverse-Grid Update: Low-Contention Gradient Aggregation↩︎

The forward kernel makes inference fit. Training gains nothing if the backward re-materializes the \([N_q, B, L_q, L_d]\) tensor and OOMs at the same batch the forward unlocked. The naive alternative, an atomicAdd of each gradient into its argmax-selected destination, serializes on hot document tokens (a frequent token can be the argmax for hundreds of queries) and doubles gradient bandwidth with FP32 atomics.

4.2.1 Closed-Form Gradients↩︎

With upstream gradient \(g_{ij} = \partial L/\partial\,\mathrm{score}(i,j)\) and the forward winner \(t^\star(i,j,s) = \arg\max_t \langle Q_{i,s}, D_{j,t}\rangle\), resolving the \(\max\) (a constant index, post-forward) and applying the chain rule gives \[\begin{align} \nabla_{Q_{i,s}} &= \textstyle\sum_j g_{ij}\, D_{j,\,t^\star(i,j,s)}, \tag{2}\\ \nabla_{D_{j,t}} &= \textstyle\sum_{(i,s):\,t^\star(i,j,s)=t} g_{ij}\, Q_{i,s}. \tag{3} \end{align}\] Eq. 2 is a gather: each query token reads exactly one \(D\) row per document, no collisions, so \(\nabla Q\) is embarrassingly parallel (one program per \((i,s)\)). Eq. 3 is a scatter: many query positions may select the same document token \(t\) (data-dependent contention), which is why \(\nabla D\) is the hard side and the inverse-grid CSR exists.

4.2.2 Runtime Inverse-Grid Construction↩︎

The forward already saves \(\mathrm{argmax} \in [0, L_d)^{[N_q\cdot B,\, L_q]}\), one index per query token. That is a source\(\to\)destination map; the backward needs the inverse: for each destination document token \((b,k)\), which query positions selected it? We invert this at runtime, inside autograd, with three primitives, bincount\(\to\)cumsum for a CSR row pointer, and a stable argsort of flattened destinations for the column indices. The CSR map costs \(O(N_q B L_q)\) for the indices plus \(O(B L_d)\) for the row pointer, far below the similarity tensor it replaces.

4.2.3 Destination-Owned Reduction↩︎

Given the CSR map, the \(\nabla D\) kernel assigns one program per output row \((b, \mathrm{tok})\); it walks that row’s source list, accumulates in an FP32 SRAM register, and writes the row once, coalesced. No two programs touch the same row, so there are no atomics. The companion \(\nabla Q\) kernel exploits the dual asymmetry (\(Q\) is shared across all \(B\) documents, so each \(\nabla Q[i]\) is written once after looping the small \(B\) axis, no CSR needed). Each ingredient is textbook; the contribution is their combination, built at runtime from a saved argmax, inside autograd, applied to a hard-selection backward, with a dispatch heuristic that falls back to atomic scatter in the low-contention regime. Benefit. Neither the CSR build nor the two row-owned kernels (Algorithm 4) ever allocate the \([N_q,B,L_q,L_d]\) similarity tensor or its gradient; peak training memory is essentially the embeddings plus the int32 argmax buffer. This is what yields the \(28\times\) memory reduction and the \(B{=}128\) batch unlock of Table 5, while keeping gradients bit-faithful to the FP32 reference.

4.2.4 Generalization to Chamfer Distance↩︎

The construction applies to any forward that emits, per source, an integer index into a destination set: argmax, argmin, top-1 attention, mixture-of-experts routing, spatial argmax (MaxPool/RoIAlign), \(k\)-NN. We instantiate it for the Chamfer distance between two point sets \(P\in\mathbb{R}^{N\times 3}\) and \(Q\in\mathbb{R}^{M\times 3}\), \[\mathrm{CD}(P,Q)=\tfrac1N\!\sum_{p}\min_{q}\lVert p-q\rVert^2 + \tfrac1M\!\sum_{q}\min_{p}\lVert q-p\rVert^2,\] the standard 3D point-cloud / shape-matching loss. It has the same structure as MaxSim with two swaps: a min-reduction over the other set (still idempotent, still rescaler-free) in place of the row-max, and squared Euclidean distance in place of the inner product. The naive form materializes the identical \([N,M]\) pairwise matrix; the fused kernel streams tiles with an online min, and the backward reuses the saved argmin (nearest-neighbour index) through the same inverse-grid CSR. Measured: gradient cosine \(1.00000\) vs.the naive autograd, \(77\) to \(152\times\) speedup, and \(100\mathrm{K}\)-point clouds that OOM the naive form run comfortably. This is direct evidence that Flash-MaxSim is a reusable hard-selection-operator pattern, not a MaxSim-specific kernel.

4.3 Quantized and Variable-Length Variants↩︎

4.3.1 INT8\(\times\)INT8 Quantization↩︎

Per-token symmetric quantization of both \(Q\) and \(D\), with dequantization fused into the kernel. This halves index storage and is \(4.7\times\) faster than a naive dequantize-then-score path at the tested shapes (not yet a net win at full ColPali scale on A100), with near-FP32 ranking fidelity (Spearman \(\rho = 0.999\) against the FP32 reference, \(100\%\) top-\(20\) overlap on tested workloads).

4.3.2 Padding-Free Variable-Length↩︎

Real corpora have ragged document lengths; batching them into a dense tensor forces padding to a common \(L_d^{\max}\), and the naive path computes (and the mask discards) all padded positions. Flash-MaxSim’s cu_seqlens variant processes only real tokens (\(\sum L_d\) work instead of \(B L_d^{\max}\)), giving up to \(5\times\) on highly ragged corpora and \(1.6\) to \(3.0\times\) at a realistic (HotpotQA-like) length distribution.

5 Experiments↩︎

5.1 Experimental Setup↩︎

NVIDIA A100-SXM4-80GB (and an H100 80GB for the forward sweep); PyTorch 2.8, FP32 matmul precision high. CUDA-event timing with warmup, median of repeated runs; per-run peak memory. The baseline is naive PyTorch (einsum + max + sum) at matched precision (FP16 inputs, FP32 accumulation). The corpus-scaling stress test (Fig. 5) is additionally run on an A100-40GB, where the OOM cliffs fall at smaller, more accessible corpus sizes.

5.2 Forward Latency and Memory↩︎

Table 1 reports forward latency: Flash-MaxSim is faster on every shape, up to \(3.9\times\) on A100 and \(4.7\times\) on H100 at ColPali scale, with the gap widening as the discarded tensor grows.

5.2.0.1 HBM traffic validates the IO bound.

Table 2 reports the HBM traffic behind Theorem 1. The naive path moves \(5\) to \(33\times\) more bytes than Flash-MaxSim (the materialized \(S\)), saturating \(\approx\!1.3\) TB/s of A100 bandwidth, i.e.it is memory-bound. Flash-MaxSim moves a constant \(0.26\) GB (inputs only) and runs near the matmul arithmetic floor (ColPali: \(1.70\) ms vs.a \(1.72\) ms floor, Table 1), confirming it has crossed from memory-bound to compute-bound.

5.2.0.2 Tile-size robustness.

Sweeping the query-chunk size (the kernel’s main tile knob) at ColPali shape, latency is flat within \(3\%\) across \(64\) to \(512\) (\(1.68\) to \(1.74\) ms) and rises to \(2.86\) ms only when chunking is disabled (\(L_q{=}1024\) in a single block), confirming both the value of the query-chunk decomposition (§4.1) and that the kernel needs no fine per-shape tile tuning.

5.2.0.3 Memory and the OOM unlock.

Figure 5 traces corpus scaling on A100-40GB for both naive precisions; Table 3 reports peak memory at ColPali shape on A100-40GB (FP16 intermediate) and A100-80GB (matched FP32, the protocol baseline). Flash-MaxSim peak memory grows linearly, tracking the document embeddings. The naive path grows as \(B L_q L_d\) and OOMs by \(B{=}20\)K on either GPU. At \(B{=}10\)K Flash-MaxSim uses \(8.2\times\) less memory than the FP16 naive path and \(16\times\) less than the matched-FP32 baseline; beyond \(B{=}20\)K only Flash-MaxSim runs.

5.3 Out-of-Core Corpus Scaling↩︎

Because Flash-MaxSim peak memory is set by a single streamed document block, a corpus that exceeds GPU memory can be scored by holding it in host RAM and streaming blocks to the GPU. Table 4 scores one ColPali query against corpora up to \(750\mathrm{K}\) documents (\(197\) GB of embeddings, \(2.3\times\) the \(85\) GB VRAM) in \(20\mathrm{K}\)-document blocks: GPU peak stays flat at \(\approx\!5.2\) GB regardless of corpus size, and steady-state throughput holds at \(\approx\!52\mathrm{K}\) documents/s, while the naive materialized path cannot fit even \(B{=}20\mathrm{K}\) on the GPU.

5.4 Training Step↩︎

For in-batch-negatives contrastive training at ColPali shape (Table 5), the naive backward materializes the \([B,B,L_q,L_d]\) tensor and its gradient (quadratic in \(B\)); Flash-MaxSim’s inverse-grid CSR backward removes both. Gradients are bit-faithful: training through Flash-MaxSim reproduces the naive loss trajectory to within \(0.001\%\) relative drift, with \(\nabla Q/\nabla D\) cosine \(1.0000\) against an FP32 reference.

At \(B{=}128\) the naive path OOMs at step 1; Flash-MaxSim

trains.

5.5 Variable-Length Scoring↩︎

On ragged corpora, where documents are padded to a common \(L_d^{\max}\) to batch them, Flash-MaxSim’s cu_seqlens variant scores only real tokens. Against the naive padded path the speedup tracks the fill ratio \(\rho = \sum_b L_d^{(b)} / (B\,L_d^{\max})\) (Table 6): \(1.6\) to \(3.0\times\) at a realistic (HotpotQA-like) distribution and up to \(5\times\) on highly ragged corpora. The result is exact (identical scores) and compounds with the dense-forward and memory gains above.

5.6 Numerical Correctness↩︎

FP32 accumulation gives max relative error \(2\times10^{-6}\) vs.an FP32 reference; top-\(20\) ranking overlap is \(100\%\) across tested shapes, confirming Proposition 1. Because the exact kernel is numerically equivalent to dense MaxSim, end-to-end retrieval quality is unchanged by construction; only the INT8\(\times\)INT8 path trades a controlled ranking fidelity (Spearman \(\rho{=}0.999\)) for speed.

6 Conclusion↩︎

Flash-MaxSim removes the similarity tensor that bottlenecks late-interaction MaxSim, in inference and in training, exactly and in one fused pass. It is faster, uses far less memory, unlocks corpus and batch sizes PyTorch cannot reach, preserves rankings exactly, and ships open-source.

6.0.0.1 Limitations.

At very small shapes the kernel-launch overhead makes Flash-MaxSim at-parity rather than faster (the launch-bound regime). The single-step training speedup is modest at matched precision; the dominant training win is memory and the batch-size unlock, not step latency. INT8\(\times\)INT8 ranking is validated on text distributions; vision (ViDoRe) end-to-end nDCG parity is not yet reported and is the primary open evaluation gap.

References↩︎