Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
10 changes: 10 additions & 0 deletions CHANGELOG.md
Original file line number Diff line number Diff line change
Expand Up @@ -4,6 +4,16 @@ All notable changes to this project will be documented here.
Format follows [Keep a Changelog](https://keepachangelog.com/en/1.1.0/) and the
project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html).

## [Unreleased]

### Removed

- The interactive kernel picker page (`docs/choose-a-kernel.html`). It predates
the API consolidation: with `maxsim` dispatching on layout and the patchers /
native PyLate & colpali-engine backends covering the framework paths, a
decision tree over many entry points no longer reflects the library. The
README "Choose a kernel" section and links are gone with it.

## [0.4.2] - 2026-06-09

### Added
Expand Down
40 changes: 11 additions & 29 deletions README.md
Original file line number Diff line number Diff line change
Expand Up @@ -16,7 +16,6 @@
---

[[How it works]](https://hcompai.github.io/late-interaction-kernels/how-it-works.html)
[[Kernel picker]](https://hcompai.github.io/late-interaction-kernels/choose-a-kernel.html)
[[Benchmarks]](docs/benchmarks.md)
[[Design]](docs/design.md)
[[Changelog]](CHANGELOG.md)
Expand Down Expand Up @@ -178,10 +177,12 @@ is asserted at `atol=1e-2` before timing.

| | Rerank /<br>inference | PyLate<br>cached-contrastive | PLAID rerank<br>vs `fast_plaid` | Fused D-head<br>(training) | FP8 vs bf16<br>(Hopper) | LateOn-Code-edge<br>e2e |
| ----------- | --------------------- | ---------------------------- | ------------------------------- | -------------------------- | ----------------------- | ----------------------- |
| **Speedup** | 1.7-16× | 5.0-6.9× | 8-23× full<br>18-51× partial | 1.5-4.5× | 1.1-1.3× | 1.00-1.06× |
| **Speedup** | 1.7-16× | 5.0-6.9× | 8-23× full<br>18-51× partial | 0.94-4.5× | 1.1-1.3× | 1.00-1.06× |

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`;
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`](docs/benchmarks.md); for how the bench scripts
themselves are organised — CLI conventions (`--only`, `--variants`),
Expand All @@ -191,16 +192,18 @@ per-script summaries, and how to run one bench, the whole sweep, or a

### Memory

The naive einsum allocates the full `[Nq · Nd · Lq · Ld]` similarity tensor as
fp32 scratch before `max(-1)`. The fused kernel never writes it: document
tiles stream through SRAM and only `[Nq, Nd]` scores come back, plus a
`[Nq · Nd, Lq]` int32 argmax buffer when training.
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.

| shape | naive scratch | fused fwd | fused fwd + bwd |
| ----------------------------------------- | ------------- | --------- | --------------- |
| `Nq=1, Nd=1k, Lq=32, Ld=300` | 183 MB | 4 KB | 128 KB |
| `Nq=1, Nd=1k, Lq=128, Ld=1024` (ColPali) | 1.0 GB | 4 KB | 512 KB |
| `Nq=16, Nd=32, Lq=32, Ld=8192` | 2.1 GB | 64 KB | 64 KB |
| `Nq=16, Nd=32, Lq=32, Ld=8192` | 2.1 GB | 2 KB | 64 KB |

The ColPali row assumes a short text query expanded to `Lq = 128`
(ColBERT-style query augmentation) against a `Ld ≈ 1024`-patch page.
Expand All @@ -216,27 +219,6 @@ 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`](docs/benchmarks.md#memory).

## Choose a kernel

<table>
<tr>
<td width="55%" valign="middle">

Not sure which entry point fits your stack? The docs site ships an interactive decision tree that narrows the public API down to the right function in four questions (stack · phase · layout · workload):

**👉 [hcompai.github.io/late-interaction-kernels/choose-a-kernel.html](https://hcompai.github.io/late-interaction-kernels/choose-a-kernel.html#choose-a-kernel)**

</td>
<td width="45%" valign="middle" align="center">

<a href="https://hcompai.github.io/late-interaction-kernels/choose-a-kernel.html#choose-a-kernel">
<img src="assets/kernel_picker_widget_preview.webp" alt="Pick a kernel · interactive decision tree" width="420">
</a>

</td>
</tr>
</table>

## API

| Symbol | What it does |
Expand Down
Binary file removed assets/kernel_picker_widget_preview.webp
Binary file not shown.
11 changes: 7 additions & 4 deletions docs/benchmarks.md
Original file line number Diff line number Diff line change
Expand Up @@ -662,17 +662,20 @@ step time, same VRAM, same numerics, one `patch_pylate()` call.

## Memory

Naive allocates `Nq · Nd · Lq · Ld · 4` bytes of fp32 scratch. Fused
keeps the `Nq · Nd` result plus — only if autograd is on — a
`Nq · Nd · Lq` int32 argmax buffer.
The naive column is the measured allocator peak. Its floor is the
`Nq · Nd · Lq · Ld · 4`-byte fp32 similarity tensor, but the measured
peak runs above that formula because the fp32 casts and broadcasts of
the operands coexist with the tensor at the peak. Fused keeps the
`Nq · Nd` result plus — only if autograd is on — a `Nq · Nd · Lq` int32
argmax buffer.


| scenario | naive scratch | fused fwd | fused fwd + argmax |
| --------------------------------- | ------------- | --------- | ------------------ |
| `Nq=1, Nd=1000, Lq=32, Ld=300` | 183 MB | 4 KB | 128 KB |
| `Nq=128, Nd=128, Lq=32, Ld=300` | 621 MB | 64 KB | 2 MB |
| `Nq=1, Nd=1000, Lq=128, Ld=1024` | 1.0 GB | 4 KB | 512 KB |
| `Nq=16, Nd=32, Lq=32, Ld=8192` | 2.1 GB | 64 KB | 64 KB |
| `Nq=16, Nd=32, Lq=32, Ld=8192` | 2.1 GB | 2 KB | 64 KB |


This is what unlocks large in-batch negatives: at the same HBM budget
Expand Down
Loading
Loading