LIK: 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 $\Dm$ tile, multiplies it against the resident $\Qm$ tile on the tensor cores, takes the row-wise max of the resulting $\Sm$ tile, and merges it into $\Mm$. After the inner loop finishes, $\Mm$ holds the true per-row maximum over the full $L_d$ axis, and the program adds $\sum \Mm$ 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 $\Sm$ is written to HBM, then read back to take the row-wise max; only the much smaller $[N_q, N_d, L_q]$ max tensor is read again for the sum. Its memory footprint is

$$\text{bytes}(\Sm) \;=\; 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 $\Sm$ is written to HBM once and read back once, all 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 $\Sm$ never exists in HBM. Only the $\text{BLOCK\_Q} \times \text{BLOCK\_D}$ slice does, and it lives in fp32 registers for the duration of one tile-matmul — SRAM holds just the $\Qm$ and $\Dm$ operand tiles feeding it.

# 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.

Per program, SRAM holds only the operand tiles $(\text{BLOCK\_Q} + \text{BLOCK\_D}) \cdot d$ in low precision — Triton double-buffers them so the next $\Dm$ tile loads while the current GEMM runs. The score tile $\Sm$ and the running max $\Mm$ live in fp32 registers, never in SRAM: tl.dot returns its accumulator straight into the register file, and the row-reduction happens there. At typical block sizes ($\text{BLOCK\_Q} = \text{BLOCK\_D} = 64$, $d = 128$), the operand tiles cost 32 KiB per pipeline stage — 64–96 KiB per program at the usual two to three stages — comfortably below the H100's 228 KiB per SM, leaving room for multiple concurrent programs and the autotuner to trade block size against occupancy.

Long-query chunking

For long query sequences ($L_q > 512$, e.g. image or document queries used on the query side), the outer loop in the pseudo-code above serialises $\sim\!8$ query-tile iterations inside one program, leaving the rest of the SM grid idle. But MaxSim is a sum over query tokens of a per-token max, so it decomposes exactly:

$$\text{score}[i, j] = \sum_{c=0}^{n_c - 1} \sum_{s \in \text{chunk}_c} \max_t \langle \Qm_{i, s},\, \Dm_{j, t} \rangle.$$

When $L_q > 512$, maxsim() reshapes $\Qm$ into $N_q n_c$ chunks of 128 tokens ($n_c = \lceil L_q / 128 \rceil$), runs the kernel over the larger batch, and sums the chunk scores back per query — plain tensor ops, so autograd flows through unchanged. Two gains at once: more programs on the grid, and the kernel always sees $L_q = 128$, pinning the autotune cache to one entry instead of one per distinct query length.

Below the $L_q = 512$ crossover the serial loop is already short and the grid already full, so the extra launches would only slow things down — short ColBERT-style queries take the un-chunked path. The same logic exempts the 4-D KD layout entirely: it already launches $N_q \cdot K$ programs, enough to fill every 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 $\Sm$ 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 $\Mm^{(0)} = -\infty$. After all tiles are consumed, $\Mm$ is the true per-row maximum and $\sum_s \Mm_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, a short chain-rule walk (Appendix D) gives:

$$\nabla_{\Qm_{i, s}} \;=\; m^q_{i,s} \cdot \sum_{j} g_{i,j} \cdot \Dm_{j,\; \operatorname{argmax}_t \langle \Qm_{i,s}, \Dm_{j,t}\rangle},$$ $$\nabla_{\Dm_{j, t}} \;=\; \sum_{(i, s)\, :\, \operatorname{argmax}[i, j, s] \,=\, t} m^q_{i,s} \cdot g_{i, j} \cdot \Qm_{i, s}.$$

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_{\Qm}$ is then embarrassingly parallel: one program per $(i, s)$ gathers $\Dm_{j, \text{argmax}}$ across $j$ and produces a single output row.

Why $\nabla_{\Dm}$ is the hard one

Picture the scatter for a fixed doc-batch index $j$. Each query-token pair $(i, s)$ contributes one row into $\nabla_{\Dm}[j, \cdot]$ at row $\operatorname{argmax}[i, j, s]$. Whether two pairs collide is a runtime property of the argmax. The kernel cannot know it statically.

This scatter is a parallel reduction, and the library implements both classic ways to organise one:

• unified — source-owned (“push”). One program per query block $(i, s)$. It already holds $\Qm_{i, s}$ for $\nabla_{\Qm}$, so it pushes its $\nabla_{\Dm}$ contribution in the same pass — both gradients from a single read of $\Qm$. Colliding writers are serialised with fp32 atomic_add, so the addition order varies run to run.
• lowmem — destination-owned (“pull”). One program per $\text{BLOCK\_D}$-row tile of $\nabla_{\Dm}[j, \cdot]$ — every output row still has exactly one writer. The program pulls the contributing $(i, s)$ pairs from the saved argmax through a one-hot matmul, accumulates in fp32 registers, and stores each row once, directly in the input dtype. No atomics, fixed order.

That last row is the entire routing logic of auto (the default). On long-query × long-doc cross-products the one-hot waste outweighs the atomic pain, so those go to unified. Everything gradient-heavy — KD / hard-negative layouts and high-contention squares ($N_q, N_d \geq 256$, $L_q \leq 64$) — goes to lowmem, which is also the path to force when bitwise-reproducible training matters.

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

Backward launch-param autotuning

The backward kernels fall into two tuning regimes. The $\nabla_{\Qm}$ kernels and the unified kernel have no block-tiling dimensions to sweep: each is one Triton program per output row streaming a single embedding vector, so they are autotuned over num_warps and num_stages only, keyed like the forward so the cache holds one entry per training regime. The lowmem $\nabla_{\Dm}$ kernel is the exception — its one-hot matmul is block-tiled just like the forward, one program per (slab, $\text{BLOCK\_D}$-row doc-tile) — so it draws from the forward's config pool and sweeps $\text{BLOCK\_Q} \times \text{BLOCK\_D}$ tile shapes alongside the launch parameters.