---
license: apache-2.0
tags:
  - kernels
  - triton
  - maxsim
  - colbert
  - late-interaction
  - sentence-transformers
---

# Late Interaction Kernels

Fused Triton kernels for MaxSim scoring, packaged for [Hugging Face Kernels](https://huggingface.co/docs/kernels) and maintained under the Sentence Transformers organization. The kernels support inference and training with token embeddings from late interaction models such as [ColBERT](https://arxiv.org/abs/2004.12832).

The kernel algorithms were developed by Aurélien Lac and Tony Wu at H Company in [Late Interaction Kernels (LIK)](https://github.com/hcompai/late-interaction-kernels). This package adapts the CUDA MaxSim implementation from [v0.4.5](https://github.com/hcompai/late-interaction-kernels/tree/4ea626dc035dc17a2b8cf5ec3a8b8c5673be8bc5) to the Sentence Transformers scoring interface. It is distributed under the [Apache 2.0 license](https://huggingface.co/kernels/sentence-transformers/late-interaction-kernels/blob/main/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:

$$
\operatorname{MaxSim}(Q, D) = \sum_{i \in \text{valid query tokens}} \max_{j \in \text{valid document tokens}} Q_i^\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:

```bash
pip install kernels triton
```

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

```python
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:

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

```python
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

| Layer | Query shape | Document shape | Output shape |
| --- | --- | --- | --- |
| `MaxSim` | `[Nq, Lq, dim]` | `[Nd, Ld, dim]` | `[Nq, Nd]` |
| `MaxSimPairwise` | `[B, Lq, dim]` | `[B, Ld, dim]` | `[B]` |

`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](https://hcompai.github.io/late-interaction-kernels/how-it-works.html) and [benchmarks](https://github.com/hcompai/late-interaction-kernels/blob/4ea626dc035dc17a2b8cf5ec3a8b8c5673be8bc5/docs/benchmarks.md) for implementation details and performance measurements.

Report issues with this package or its use in Sentence Transformers in the [Hub discussions](https://huggingface.co/kernels/sentence-transformers/late-interaction-kernels/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](https://github.com/hcompai/late-interaction-kernels#citation):

```bibtex
@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},
}
```