Scoring every query–document token pair without materialising the table

1The MaxSim score

MaxSim is the scoring function ColBERT-style late-interaction models use. The inputs are two batches of token-level embeddings:

• $\Qm \in \mathbb{R}^{N_q \times L_q \times d}$: $N_q$ queries, each with $L_q$ token vectors of dimension $d$.
• $\Dm \in \mathbb{R}^{N_d \times L_d \times d}$: same shape, for documents.

For every (query $i$, document $j$) pair, the score is:

$$\operatorname{score}[i, j] \;=\; \sum_{s=1}^{L_q} \; \max_{t=1}^{L_d} \; \langle \Qm_{i, s}, \Dm_{j, t} \rangle.$$

For each query token, find the document token it matches best (highest inner product), then sum those best matches across all query tokens. The result is a score matrix in $\mathbb{R}^{N_q \times N_d}$. Optional boolean masks drop padding tokens from the sum and from the inner max.

2Programs and the launch grid

A GPU kernel is a small function the device runs in many parallel instances. Each instance is a program in Triton's vocabulary, or a thread block in CUDA's. Writing a kernel comes down to two design choices: how many programs to launch, and what one program computes.

The MaxSim forward uses a grid of $N_q \cdot N_d$ programs, one per query–document pair. Program $(i, j)$ pulls the embeddings of query $i$ and document $j$ from HBM, accumulates a single fp32 scalar, and writes it to position $(i, j)$ of the score matrix. All $N_q \cdot N_d$ programs run concurrently on the GPU's streaming multiprocessors, with the hardware scheduling them onto whatever SMs are free.

3Inside one program

Fix a single program at $(i, j)$. It has the $i$-th query and the $j$-th document available in HBM. The most direct implementation of its scalar output looks like

$$\Sm \;=\; \Qm_{i,\,\cdot}\, \Dm_{j,\,\cdot}^{\top} \;\in\; \mathbb{R}^{L_q \times L_d}, \qquad \operatorname{score}[i,j] \;=\; \sum_{s} \max_{t} \Sm_{s, t}.$$

Building $\Sm$ even for a single pair wastes most of the values: every row of $\Sm$ contributes only one number to the final score, the row maximum. The kernel never builds it. Instead the program walks the output in tiles. Pick block sizes $B_q$ and $B_d$ (usually 32 or 64). The program loops over $\lceil L_q / B_q \rceil$ query tiles and, inside each, over $\lceil L_d / B_d \rceil$ document tiles. The largest thing that ever exists is one $B_q \times B_d$ slice of $\Sm$, held in registers for the duration of a single tile-matmul.

Each inner-loop iteration loads a fresh $D$ tile, multiplies it against the resident $Q$ tile on the tensor cores, takes the row-wise max of the resulting $S$ tile, and merges it into $m$. After the inner loop finishes, $m$ holds the true per-row maximum over the full $L_d$ axis, and the program adds $\sum m$ to its scalar accumulator. After the outer loop finishes, that accumulator is written to HBM. One pair, one scalar, one store.

4Why the naive form is bandwidth-bound

The textbook implementation of MaxSim, the one inside PyLate's reference code, separates the two reductions:

S = einsum("nsd, mtd -> nmst", Q, D)   # [Nq, Nd, Lq, Ld]  ← lives in HBM
M = S.max(dim=-1).values               # [Nq, Nd, Lq]      ← max over t
scores = M.sum(dim=-1)                 # [Nq, Nd]          ← sum over s

The full similarity tensor $S$ is written to HBM, then read back to take the row-wise max, then read again for the sum. Its memory footprint is

$$\text{bytes}(S) \;=\; N_q \cdot N_d \cdot L_q \cdot L_d \cdot \text{sizeof(float)}.$$

Scaling is quadratic in the sequence length. At ColPali shapes ($N_q = N_d = 64$, $L_q = L_d = 1024$, fp32) the tensor is $16$ GB; in fp16 it is $8$ GB. Training builds three or four such intermediates per step. The chart shows the same number plotted as a function of the sequence length:

Beyond the memory pressure, the naive version is bandwidth-bound: each element of $S$ is written to HBM once and read twice for a single multiply-add per element. The arithmetic intensity is far below the H100 ridge point, so the SMs spend most of their time waiting on memory. Appendix B works the FLOP-per-byte numbers out and places both implementations on the roofline.

5Fused forward: tile, stream, accumulate

Spelling out the loop nest sketched in section 3: one Triton program per $(q\text{-batch}, d\text{-batch})$ pair, document tiles of $\text{BLOCK\_D}$ rows nested inside query tiles of $\text{BLOCK\_Q}$ rows. The full similarity tensor $S$ never exists in HBM. Only the $\text{BLOCK\_Q} \times \text{BLOCK\_D}$ slice does, and it lives in SRAM for the duration of one tile-matmul.

# one program per (q_batch, d_batch)
score ← 0
for q_start in range(0, Lq, BLOCK_Q):
    Q_tile ← Q[q_batch, q_start : q_start + BLOCK_Q]      # SRAM
    m     ← −∞                                            # [BLOCK_Q] in registers
    for d_start in range(0, Ld, BLOCK_D):
        D_tile ← D[d_batch, d_start : d_start + BLOCK_D]  # SRAM
        S_tile  = Q_tile @ D_tileᵀ                        # tensor cores, fp32 acc
        S_tile ← mask(S_tile, d_active, −∞)               # fused-in masking
        m      ← maximum(m, rowmax(S_tile))               # online max
    score += sum(m)                                       # contribution from this Q tile
scores[q_batch, d_batch] ← score                           # only this scalar is written

The structure is the same outer-product tiling that FlashAttention uses, with one simplification: where FlashAttention has to maintain a running max and a running sum of exponentials for softmax, MaxSim only has a running max. The recurrence is exact and there is nothing to rescale.

Peak SRAM per program is $(\text{BLOCK\_Q} + \text{BLOCK\_D}) \cdot d$ for the operand tiles plus $\text{BLOCK\_Q} \cdot \text{BLOCK\_D}$ for the score tile, in low precision, with a $[\text{BLOCK\_Q}]$ fp32 accumulator on top (roughly 40 KiB at typical block sizes), leaving room for four concurrent programs per H100 SM.

Step through the algorithm

Below is a toy execution with $L_q = 4$, $L_d = 12$, $d = 4$, $\text{BLOCK\_Q} = 4$ (one outer iteration) and $\text{BLOCK\_D} = 4$ (three inner iterations). Use the controls to step through; the counters on the right track HBM bytes moved versus the naive baseline.

6The online-max recurrence

To eliminate the score tile entirely, the inner loop replaces every value of $S$ with its contribution to the per-row maximum, which we update incrementally. After consuming the $k$-th tile, the running max obeys

$$\Mm^{(k)} \;=\; \max\!\bigl(\Mm^{(k-1)},\;\; \operatorname{rowmax}(\Sm^{(k)})\bigr),$$

with $m^{(0)} = -\infty$. After all tiles are consumed, $m$ is the true per-row maximum and $\sum_s m_s$ is the score contribution for the current $\text{BLOCK\_Q}$. Both quantities are computed in fp32 in registers; the GEMM uses bf16/fp16 operands with fp32 accumulation, matching the tensor-core native path.

Worked example

Three doc-tiles, one query row. The argmax index is recorded for backward.

Masked positions are written as $-\infty$ before the reduction, so they cannot influence the argmax even when scores would otherwise be negative. This is stricter than PyLate's reference, which post-multiplies by a $0/1$ mask, and matches the masking discipline used by flash-maxsim and FlashAttention.

7Backward: argmax-only gradients

$\max$ is sub-differentiable: only the argmax position carries a gradient, so the backward pass is simpler than for softmax. With $g \in \mathbb{R}^{N_q \times N_d}$ the upstream gradient on the score:

$$\nabla_Q \, Q_{i, s} \;=\; q\text{-}\text{mask}_{i,s} \cdot \sum_{j} g_{i,j} \cdot D_{j,\; \operatorname{argmax}_t \langle Q_{i,s}, D_{j,t}\rangle}.$$

$\nabla_D$ has the symmetric form, summed over the $(i, s)$ pairs whose argmax falls on the same $(j, t)$. To avoid recomputing the forward, the forward kernel optionally writes an $[N_q \cdot N_d, L_q]$ int32 buffer of argmax indices (4 MB for a typical training batch).

$\nabla_Q$ is embarrassingly parallel: one program per $(i, s)$ gathers $D_{j, \text{argmax}}$ across $j$ and produces a single output row. $\nabla_D$ has output contention: many $(i, s)$ pairs can land on the same $(j, t)$. The library offers three reduction strategies:

All three paths use Triton's stable argmax (lowest-index tie-break) on the forward, so only the $\nabla_D$ reduction order distinguishes them numerically.

8Performance summary

Measured on a single H100 80 GB SXM with bf16/fp16 inputs and fp32 accumulators, 50-iteration median; baselines are the same operation in plain PyTorch. Full shapes and reproduction commands are in docs/benchmarks.md.

On full encoder training the MaxSim step is dominated by the transformer forward/backward, so the kernel is effectively a free drop-in. Everywhere MaxSim is not negligible (inference, reranking, long documents, knowledge distillation, compressed indices, small encoders), the fused path dominates.

9Variants in the library

The streaming-and-never-storing structure is the entire optimisation. The rest of the library applies the same idea to a few adjacent workloads:

• Variable-length / packed input (maxsim_varlen). Real corpora have ragged lengths; padding to $L_d^{\max}$ wastes roughly half the FLOPs. The packed kernel consumes $[\text{total\_tokens}, d]$ tensors with FlashAttention-style cu_seqlens offsets and runs the same tiling.
• Fused $D$-side projection (maxsim_from_hidden). For corpora stored as ModernBERT hidden states, the kernel folds $\texttt{Linear} \rightarrow \texttt{L2-normalize} \rightarrow \texttt{MaxSim}$ into a single pass. The projected embedding tensor is never written to HBM. The training backward gathers only the winning positions and recomputes the projection locally.
• PLAID / ColBERTv2 (maxsim_residual, maxsim_residual_varlen). The doc embeddings live on disk as $(\text{centroid index}, \text{quantised residual})$. The kernel decompresses on the fly in SRAM, optionally L2-normalises, and runs MaxSim, all in one program with no dense decompressed tensor ever materialised.
• FP8 MaxSim (maxsim_inference_fp8). On Hopper, the GEMM uses the FP8 tensor cores at twice the bf16 throughput; the reduction stays in fp32 to preserve range.

Different inputs, same recipe: keep large intermediates on-chip, run the reductions in registers, and ship only the final scores to HBM.