Late Interaction Kernels

Fused Triton kernels for MaxSim scoring, packaged for Hugging Face Kernels and maintained under the Sentence Transformers organization. The kernels support inference and training with token embeddings from late interaction models such as ColBERT.

The kernel algorithms were developed by Aurélien Lac and Tony Wu at H Company in Late Interaction Kernels (LIK). This package adapts the CUDA MaxSim implementation from v0.4.5 to the Sentence Transformers scoring interface. It is distributed under the Apache 2.0 license.

What does it compute?

For each query token, MaxSim finds the document token with the highest dot product, then sums those maxima over the query:

$$\mathrm{MaxSim}(Q,D)=\sum _{i\in \text{valid query tokens}}\underset{j\in \text{valid document tokens}}{\max }{Q}_i^{\mathrm{\top }}{D}_j\backslash operatorname\{MaxSim\}(Q,\; D)\; =\; \backslash sum\_\{i\; \backslash in\; \backslash text\{valid\; query\; tokens\}\}\; \backslash max\_\{j\; \backslash in\; \backslash text\{valid\; document\; tokens\}\}\; Q\_i^\backslash top\; D\_j$$

The fused implementation computes similarities and reductions without materializing the full [Nq, Nd, Lq, Ld] token similarity tensor in GPU memory. This reduces intermediate memory use and can accelerate scoring, especially for larger batches or longer documents. Both scoring layers support autograd, so gradients flow into query and document embeddings.

Quickstart

Use a Linux environment with CUDA-enabled PyTorch, Triton, and a supported NVIDIA GPU (compute capability 7.5 or newer). Install the following alongside your CUDA-enabled PyTorch installation:

pip install kernels triton

The following example scores two queries against four documents, then scores two matching pairs:

import torch
import torch.nn.functional as F
from kernels import get_kernel

kernel = get_kernel(
    "sentence-transformers/late-interaction-kernels",
    version=0,
    trust_remote_code=True,
)

queries = F.normalize(torch.randn(2, 32, 128, device="cuda"), dim=-1).half().requires_grad_()
documents = F.normalize(torch.randn(4, 180, 128, device="cuda"), dim=-1).half().requires_grad_()
query_mask = torch.ones(2, 32, dtype=torch.bool, device="cuda")
document_mask = torch.ones(4, 180, dtype=torch.bool, device="cuda")
document_mask[:, 150:] = False

scores = kernel.layers.MaxSim()(queries, documents, query_mask, document_mask)
paired_scores = kernel.layers.MaxSimPairwise()(
    queries, documents[:2], query_mask, document_mask[:2]
)

print(scores.shape)  # torch.Size([2, 4])
print(paired_scores.shape)  # torch.Size([2])
scores.sum().backward()

trust_remote_code=True allows the loader to execute the Python code in this kernel repository. The first call for a new input configuration may take longer while Triton compiles and autotunes the kernel.

Sentence Transformers

With Sentence Transformers 6.3.0 or newer, set use_kernels=True when loading a MultiVectorEncoder to enable these MaxSim kernels automatically on supported CUDA devices:

pip install "sentence-transformers[kernels]>=6.3.0" triton

from sentence_transformers import MultiVectorEncoder

model = MultiVectorEncoder("lightonai/LateOn", device="cuda", use_kernels=True)
queries = model.encode_query(["What is the capital of France?"])
documents = model.encode_document(["Paris is the capital of France.", "Berlin is the capital of Germany."])
scores = model.similarity(queries, documents)

Scoring interface

MaxSim scores every query against every document, for example for in-batch training. MaxSimPairwise scores corresponding pairs, for example for reranking, without computing the full cross product.

• Inputs must share the embedding dimension, device, and dtype.
• Masks have shape [batch_size, sequence_length]. Use boolean masks with True for valid tokens and False for padding, on the same device as the embeddings. Pass None for either mask to include all tokens.
• Query masks exclude tokens from the sum. Document masks exclude tokens from the maximum. Documents with no valid tokens receive a score of -1e9.
• The layers use dot products and do not normalize embeddings internally. Normalize each token embedding beforehand if your model uses cosine similarity.
• Scores are returned in torch.float32.

Precision and compatibility

Supports float16, bfloat16, and float32 inputs. The accelerated path uses half precision with float32 accumulation, so scores and gradients may differ slightly from full precision PyTorch scoring.

The published kernel targets CUDA. Inputs that cannot use the fused kernel fall back to PyTorch scoring, as does execution inside torch.compile.

Performance and support

See the upstream algorithm walkthrough and benchmarks for implementation details and performance measurements.

Report issues with this package or its use in Sentence Transformers in the Hub discussions. Include your GPU, PyTorch and Triton versions, input shapes and dtypes, and a minimal reproducer.

Citation

Please credit the original kernel authors using the upstream 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},
}