diff --git a/CHANGELOG.md b/CHANGELOG.md
index 38e6dcd..338bd2a 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -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
diff --git a/README.md b/README.md
index c5b62c1..362afef 100644
--- a/README.md
+++ b/README.md
@@ -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)
@@ -178,10 +177,12 @@ is asserted at `atol=1e-2` before timing.
| | Rerank / inference | PyLate cached-contrastive | PLAID rerank vs `fast_plaid` | Fused D-head (training) | FP8 vs bf16 (Hopper) | LateOn-Code-edge e2e |
| ----------- | --------------------- | ---------------------------- | ------------------------------- | -------------------------- | ----------------------- | ----------------------- |
-| **Speedup** | 1.7-16× | 5.0-6.9× | 8-23× full 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 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`),
@@ -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.
@@ -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
-
-
-
-
-
-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)**
-
-
- The library exposes everything from a one-line PyLate drop-in to raw fused
- Triton kernels for PLAID-compressed corpora. Answer four questions about
- your stack and workload, and the widget below narrows the surface down to
- the function you should call.
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
diff --git a/docs/design.md b/docs/design.md
index a6824dd..161b2ae 100644
--- a/docs/design.md
+++ b/docs/design.md
@@ -75,8 +75,9 @@ Two benefits:
`static_range(Lq / BLOCK_Q)` loop inside each program. Chunking launches
more, shorter programs that fill the SM array concurrently.
* **Autotune cache collapse.** The kernel always sees `Lq == 128`, so
- all long queries share one autotune entry (two with tail-padding)
- instead of one per `Lq` bucket.
+ all long queries share one autotune entry — tail-padded and
+ exact-multiple lengths included, since the mask flag is out of the
+ autotune key — instead of one per `Lq` bucket.
The split fires only when `Lq > 512`; shorter queries fall through to the
direct kernel call unchanged. Chunking is cross-product-only — the KD /
@@ -205,11 +206,14 @@ the key — they're constexpr/runtime toggles that change codegen but not the
winning tile, so keeping them only multiplied the sweep count. Configs are
pruned by:
-* shared-memory budget (≤ 224 KiB on H100, ≤ 164 KiB on A100),
+* shared-memory budget (≤ 220 KiB on H100, ≤ 156 KiB on A100 — the raw
+ 228 / 164 KiB per SM minus an 8 KiB reserve for Triton scratch),
* `BLOCK_Q > 2 · Lq` (block bigger than the problem).
-Hopper shortlist enables `num_consumer_groups` warp specialization
-(Triton ≥ 3.2, FA-3 style); fallback is transparent on older Triton.
+All configs are plain `(BLOCK_Q, BLOCK_D, num_warps, num_stages)` tuples
+with no warp-specialization kwargs; the old `num_consumer_groups` Hopper
+entries were dropped when Triton 3.5 removed the API (see the v0.2.0
+changelog).
With long-query chunking (see `## Forward → Long-query chunking`) long
queries always present `Lq = 128` to the kernel, collapsing the forward
diff --git a/docs/how-it-works.html b/docs/how-it-works.html
index ea82818..8ed96fb 100644
--- a/docs/how-it-works.html
+++ b/docs/how-it-works.html
@@ -916,8 +916,9 @@
1The MaxSim score
The problem to solve.
Computing this naively builds the full $[N_q, N_d, L_q, L_d]$ similarity
- tensor first, then reduces. At ColPali scale ($L_q = L_d = 1024$, batch
- of 64 pairs) that tensor is 8 GB in fp16. Each element is written once
+ tensor first, then reduces. At ColPali scale ($L_q = L_d = 1024$, a
+ batch of 64 queries × 64 documents — 4,096 scored pairs) that tensor is
+ 8 GB in fp16. Each element is written once
and read back once, so the GPU spends most of its time on HBM traffic,
not compute. The kernels in this library compute the same result without ever
materialising that tensor, by streaming $B_q \times B_d$ tiles through
@@ -1200,58 +1201,58 @@
S in HBM, fp16, $N_q = N_d = 64$fused kernel — only the $[N_q, N_d]$ score matrix
+ fp32 doubles every point; an 80 GB H100 runs out of memory near $L \approx 3125$ even in fp16
@@ -1275,8 +1276,9 @@
5Fused forward: tile, stream, accumulate
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 SRAM
- for the duration of one tile-matmul.
+ $\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)
@@ -1397,7 +1399,7 @@
HBM traffic (one program)
Q (active)D (active)
- S tile · ephemeral, SRAM only
+ S tile · ephemeral, registers onlyrunning max mconsumed
@@ -1431,8 +1433,9 @@
6The online-max recurrence
softmax has to track $\Mm$ and the running sum of exponentials
$\ell$, and rescale both whenever the running max changes: the rescalers
are the $\exp(\Mm^{(k-1)} - \Mm^{(k)})$ factors in FA-1. Online max has no such
- rescaler: $\max$ is idempotent, so the update is just an elementwise max
- and the result is identical to the offline computation.
+ rescaler: $\max$ is associative and commutative, so folding the tiles in
+ one by one — in any order — gives exactly the same result as the offline
+ computation.
Worked example
@@ -1544,10 +1547,12 @@
Why $\nabla_{\Dm}$ is the hard one
lowmem — destination-owned (“pull”).
- One program per output row $(j, t)$ — that row's only writer. It
- pulls the contributing $(i, s)$ pairs from the saved argmax
- through a one-hot matmul, accumulates in fp32 registers, and stores
- once, directly in the input dtype. No atomics, fixed order.
+ 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.
@@ -1616,12 +1621,12 @@
Why $\nabla_{\Dm}$ is the hard one
(i=2, s=2)
- → pulled by the (j, t=12) owner
+ → pulled by the (j, t=12…15) owner
- program owns (j, t=5)
+ program owns tile (j, t=4…7)Σ in fp32 registersvia one-hot matmul
@@ -1632,7 +1637,7 @@
- The backward kernels 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 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.
@@ -1713,7 +1725,7 @@
AGPU memory hierarchy
Level
Size
Bandwidth
Where it lives
-
Registers
256 KB / SM
~20 TB/s
Inside each thread
+
Registers
256 KB / SM
~80 TB/s
Inside each thread
SRAM (shared memory)
228 KB / SM (H100)
~19 TB/s
On-chip, per SM
L2 cache
50 MB (H100)
~5 TB/s
On-chip, chip-wide
HBM
80 GB (H100)
~3.3 TB/s
Off-chip, on the package
@@ -1833,18 +1845,18 @@
BArithmetic intensity
CThe constexpr bargain
- A tl.constexpr kernel argument is baked into the compiled
- binary. That is a bargain, not a feature: the compiler can specialise
- and unroll code for that exact value, but every new value it meets
- compiles a brand-new kernel — a multi-second stall — and adds an entry
- to the autotune cache. So every shape-like argument poses the same
- question: does the compiler actually do anything with this
- value?
+ A tl.constexpr kernel argument is a value the compiler
+ sees at compile time, so it can generate code tailored to exactly
+ that value. The catch: every new value triggers a fresh
+ compile (a multi-second stall) and adds another autotune-cache entry.
+ That is the bargain, and it is only worth taking when the answer to
+ one question is yes: does knowing the value let the compiler
+ generate better code?
- The forward kernel answers it twice, on two loops that look almost
- identical:
+ The forward kernel answers it both ways, on two loops that look
+ almost identical:
for q_start in tl.static_range(0, Lq, BLOCK_Q): # unrolled at compile time → Lq is constexpr
@@ -1852,33 +1864,37 @@
CThe constexpr bargain
for d_start in range(0, Ld, BLOCK_D): # plain runtime loop → Ld is a normal argument
- For $L_d$ the answer is no. It bounds the long inner loop, which the
- compiler pipelines rather than unrolls, so specialising on it would
- make $L_d = 300$ and $L_d = 301$ different binaries with identical
- machine code — one recompile per distinct document length in a ragged
- corpus, for nothing. $L_d$ stays a plain runtime argument.
+ For $L_d$: no. The inner loop is long, so the compiler keeps it as a
+ loop and pipelines it, loading the next $\Dm$ tile while the current
+ one is being multiplied. The machine code is identical whether
+ $L_d$ is 300 or 301; specialising would just recompile once per
+ distinct document length in a ragged corpus, for nothing. $L_d$ stays
+ a plain runtime argument.
+
+
+
+ For $L_q$: yes. The outer loop runs just
+ $n = \lceil L_q / \text{BLOCK\_Q} \rceil$ iterations, so the compiler
+ unrolls it: it deletes the loop and pastes $n$ copies of the
+ body back-to-back — straight-line code, no counter, no branch. It can
+ only do that if it knows $n$ at compile time, and $n$ comes from
+ $L_q$. That is what makes $L_q$ constexpr.
- For $L_q$ the answer is yes. The outer loop runs only
- $n = \lceil L_q / \text{BLOCK\_Q} \rceil$ iterations — a handful of
- query tiles — so it is worth unrolling: the compiler deletes
- the loop and pastes $n$ consecutive copies of its body (the
- load-tile / matmul / row-max sequence) as straight-line code, no
- counter, no branch. It can only do that if it knows $n$ at compile
- time, and $n$ comes from $L_q$ — that is what forces $L_q$ to be
- constexpr. The price would be one recompile per distinct query
- length, and per-batch lengths float step to step. So the kernel
- shrinks the domain instead of refusing the deal: $L_q$ is rounded up
- to the next power of two and $\Qm$ zero-padded, with the query mask
- covering the fake tokens exactly — at most nine values between 16 and
- 4096 can ever exist.
+ One problem remains: query lengths float from batch to batch, and
+ each distinct $L_q$ would be its own compile. So the kernel shrinks
+ the domain. $L_q$ is rounded up to the next power of two, $\Qm$ is
+ zero-padded, and the query mask hides the padding — at most nine
+ compiled values between 16 and 4096 in normal use. (Past 4096, $L_q$
+ passes through unbucketed, with a one-shot warning pointing
+ long-context callers at maxsim_varlen.)
-
-
-