late-interaction-kernels

---

[How it works] [Benchmarks] [Design] [Changelog]

The full algorithmic walkthrough (tiling, online max, the backward pass) with step-through animations and benchmark plots lives on the docs site: 👉 hcompai.github.io/late-interaction-kernels

Introduction

late-interaction-kernels provides fused Triton and Metal kernels for MaxSim, the late-interaction scoring at the heart of ColBERT, ColPali, ModernColBERT, LateOn and ColBERTv2. They're numerically identical to plain PyTorch, but fuse the similarity matrix, max-reduction and (optional) L2-normalisation into a single launch, so the full [Nq, Nd, Lq, Ld] score tensor never lands in HBM.

PyLate and colpali-engine support them natively: install the extra and their auto dispatch picks the kernels up, no code change. You can also call them directly: a stateless MaxSimScorer module for custom training loops, or function-level entry points (maxsim, maxsim_varlen, maxsim_padded, ...) for everything else.

Install

pip install late-interaction-kernels

Quickstart

Score directly (maxsim / maxsim_pairs)

maxsim is the lowest-level public entry point: autograd-aware, mask-aware, and dispatches on D.dim(), so one call covers in-batch and knowledge-distillation layouts in a single fused launch. The argmax buffer for the backward is skipped automatically when neither input requires grad, so this is the inference path too.

from late_interaction_kernels import maxsim, maxsim_pairs

# in-batch:  Q[Nq, Lq, d] × D[Nd, Ld, d]    → [Nq, Nd]
scores = maxsim(Q, D, q_mask=q_mask, d_mask=d_mask, normalize=True)

# KD / hard-negative:  D is 4D [Nq, K, Ld, d]  → [Nq, K]   (one launch, no Python loop)
scores = maxsim(Q, D_kd, q_mask=q_mask, d_mask=d_mask_kd)

# pairwise (diagonal):  Q[B, Lq, d] × D[B, Ld, d]  → [B]
scores = maxsim_pairs(Q, D, q_mask=q_mask, d_mask=d_mask)

PyLate & colpali-engine

Both ship a native LIK backend: install the extra and their auto dispatch picks it up, no code change (force it with PYLATE_SCORES_BACKEND=lik / COLPALI_SCORES_BACKEND=lik). On older versions the patch_* drop-ins route scoring + loss through the fused kernel at import time (LIK_DISABLE=1 falls back; deprecated no-ops once native support is present).

pip install "pylate[lik]"

import late_interaction_kernels as lik
lik.patch_pylate()

pip install "colpali-engine[lik]"

import late_interaction_kernels as lik
lik.patch_colpali_engine()

Top-k retrieval

Score Q against a large corpus and return the top-k per query without materialising the full [Nq, Nd] matrix. chunk= streams documents in tiles so peak HBM stays bounded.

from late_interaction_kernels import retrieve

scores, indices = retrieve(Q, D, top_k=100, chunk=4096)
# both [Nq, 100]; chunk= bounds peak HBM at Nq * (chunk + top_k)

from late_interaction_kernels.plaid import maxsim_residual_varlen

scores = maxsim_residual_varlen(
    Q, codes_flat, residuals_flat, cu_seqlens_d,
    centroids=centroids, bucket_weights=bucket_weights,
    nbits=2, normalize=True,
)  # [Nd] fp32; one kernel does decompress + L2-normalize + MaxSim

from late_interaction_kernels import MaxSimScorer

scorer = MaxSimScorer(normalize=True)                # nn.Module, no parameters
scores = scorer(Q, D, q_mask=q_mask, d_mask=d_mask)  # [Nq, Nd] fp32
scores.mean().backward()

Benchmarks

1×H100 80GB SXM, bf16 inputs / fp32 accumulator, 50-iter median. All speedups are measured at matched numerics: every baseline runs the einsum with an fp32 accumulator (same as the fused kernel), and parity is asserted at atol=1e-2 before timing.

Speed

Rerank is vs both the eager fp32-accumulator path and torch.compile; PLAID rerank includes top-k; the fused D-head win grows with Nd · Ld (the two smallest LateOn shapes are 0.94-0.95×, i.e. slightly slower; ≥1.4× from Nd=128, Ld=1024 up — every ColBERT/ColPali-scale shape); FP8 is at Ld ≥ 256. Full tables and reproduction commands live in docs/benchmarks.md; for how the bench scripts themselves are organised — CLI conventions (--only, --variants), per-script summaries, and how to run one bench, the whole sweep, or a RUN_ONLY-filtered subset on a SkyPilot cluster — see benchmarks/README.md.

Memory

The naive einsum materialises the full [Nq · Nd · Lq · Ld] similarity tensor in fp32 before max(-1); its column reports the measured allocator peak, which runs above the similarity tensor alone because the fp32-cast operand copies coexist with it. The fused kernel never writes any of that: document tiles stream through SRAM and only [Nq, Nd] scores come back, plus a [Nq · Nd, Lq] int32 argmax buffer when training.

The ColPali row assumes a short text query expanded to Lq = 128 (ColBERT-style query augmentation) against a Ld ≈ 1024-patch page.

This runs long-context shapes (Ld ≥ 8k) that OOM the naive path, and fits ~5–10× more in-batch negatives at a fixed HBM budget. In real ColQwen2 training (80 GB H100, LoRA + grad-ckpt, vidore/colpali_train_set) vanilla colpali-engine OOMs at batch=128 where the MaxSim op holds 7.8 GiB; the fused kernel holds 61 MiB and doubles the batch ceiling at the same step time. The backward keeps the discipline: auto routes gradient-heavy shapes to lowmem, writing grad_Q / grad_D in the input dtype (no full-size fp32 buffer, no atomics, deterministic) for roughly half the backward peak, e.g. a B256 × 16-neg ColPali step from 4.3 GB to 2.2 GB. Full tables in docs/benchmarks.md.

API

Other kernels are in submodules: padded, score_pairs, fused_head, plaid, fp8, reference. See docs/design.md for details on every kernel, the autograd graph and the backward variants.

Development

git clone https://github.com/hcompai/late-interaction-kernels
cd late-interaction-kernels
uv sync --extra dev --extra pylate --extra torch-cuda   # GPU dev; use --extra torch-cpu on CPU-only boxes
uv run pytest -q                                        # CUDA tests auto-skip without a GPU
uv run ruff check . && uv run ruff format --check .

[!NOTE] Pick exactly one of --extra torch-cuda (pulls torch from the CUDA index — cu124) or --extra torch-cpu (CPU-only wheel, what CI uses). The two are declared as conflicting in pyproject.toml so the lockfile resolves cleanly for both. On macOS, --extra torch-cpu falls back to PyPI's default (MPS-capable) wheel automatically.

See CONTRIBUTING.md for the contribution workflow, including how GPU tests run.

Related projects

Citation

@software{late_interaction_kernels_2026,
  author  = {Lac, Aurélien and Wu, Tony},
  title   = {{late-interaction-kernels}: Fused Triton kernels for late-interaction scoring},
  year    = {2026},
  url     = {https://github.com/hcompai/late-interaction-kernels},
}