TileMaxSim

GPU Triton kernels for MaxSim scoring in multi-vector retrieval (ColBERT-style late interaction), and the measurements behind the paper "Fused MaxSim on an H100: Two Roofline Regimes, a Head-to-Head, and Why Product Quantization Loses".

A naive MaxSim implementation materializes the full N_q x N_d query-by-document similarity matrix in HBM and then reduces it. These kernels do not: document embeddings stream through SRAM, per-query-token running maxima stay in registers, and each embedding is read from HBM once when the query fits in one tile.

This repository is the artifact for the paper: source, configs, scripts and the result files each number comes from. It does not contain the paper source.

Read this before quoting any number

The paper was substantially revised in September 2026 and several figures from the first version were withdrawn. The arXiv v2 "Comments" field lists all of them; the three that matter most to anyone reading this code:

1. This is not the fastest fused MaxSim kernel, and that is measured. Two others were installable with pip before our first posting: flash-maxsim (Pony et al., PyPI 0.1.0 on 2026-03-19, 0.3.0 with a split-d path on 2026-06-11) and late-interaction-kernels (Lac and Wu, PyPI from 2026-05-02, shipped as a backend in PyLate and colpali-engine). Measured head to head over 28 shapes with every row checked against an exhaustive FP32 reference, this kernel is faster on 9 and slower on 19 (median 0.952x, range 0.343x to 1.399x). In the ColBERT regime all three land within a few percent of each other and all reach 90 to 94% of the 3,350 GB/s theoretical peak. The remaining large losses all have one cause: BLOCK_Nq caps at 64, so a query longer than that makes multiple passes over the documents. Artifact: experiment_results/remeasure_2026_09/head2head/head2head_merged.json.

2. Dimension tiling is not a distinguishing feature. It is the standard split-K matmul pattern, and Pony et al. shipped it for this operator before our first posting. Earlier versions of this README and of the paper called it a distinguishing contribution; that claim is withdrawn.

3. Fused product quantization is a negative result, not a feature. The fused ADC path moves 15.9x fewer bytes and is still 3.0x slower than the exact FP16 tensor-core path, because it is limited by the number of dependent gathers rather than by bytes moved: 1.23% of peak DRAM throughput at 98.2% occupancy. Rewriting it in CUDA with an explicitly staged table recovers 4.7x and still loses. Read this before building on src/flash_pqsim_kernel.py.

Also withdrawn: the "~80% of peak" headline, which used an "achievable" 3,024 GB/s denominator that the kernel itself exceeds (it reaches 3,147 GB/s), so percentages are now given against the 3,350 GB/s theoretical peak only; a 220x speedup measured against baselines that upcast FP16 to FP32; and a 469x figure against WARP that compared incommensurable units.

Two roofline regimes

The operator does not sit on one side of the roofline. Fused arithmetic intensity is Nq*Nd/(Nq+Nd) against an H100 dense-FP16 ridge of 295.4 FLOP per byte (989.5 TFLOP/s over 3.35 TB/s; do not use 1979 TFLOP/s, which is the 2:4-sparsity figure). So percent of peak bandwidth is only a meaningful metric below the crossover. Predicted crossover Nq = Nd*r/(Nd-r) is 698, 415 and 345 at Nd of 512, 1024 and 2048, and the measured brackets are (512, 1024), (256, 512) and (256, 512): the prediction falls inside all three. Artifacts under experiment_results/remeasure_2026_09/regimes/.

Layout

src/
  flash_maxsim_v2.py          # V2 and V2-MQ (multi-query) kernels; V2-MQ is the one reported
  flash_maxsim_kernel.py      # V1 per-query-token kernel and PyTorch baselines
  flash_pqsim_kernel.py       # fused PQ/ADC scoring (negative result, see above)
  pqsim_cuda.cu               # hand-written CUDA ADC kernel with a staged table
  head2head_2026_09.py        # head-to-head vs flash-maxsim and late-interaction-kernels
  validate_dim_dispatch.py    # correctness and speed of the embedding-dimension dispatch
  regenerate_all_tables.py    # regenerates every paper table from one build
  remeasure_2026_09.py        # int64 validation, bandwidth sweep, dtype-matched baselines
  beir_evaluation.py          # BEIR quality verification
  e2e_retrieval_quality.py    # MS MARCO re-ranking equivalence check
baselines/                    # WARP and PLAID comparison harnesses
scripts/                      # reproduction and profiling drivers
experiment_results/           # results; see PROVENANCE.md for which file backs which number

Module filenames still say flash_maxsim because many scripts import them. The project is TileMaxSim; renaming the modules is pending.

Requirements

• An NVIDIA GPU. Every number in the current paper was measured on a single H100 80GB HBM3 with CUDA 12.8.
• Python 3.13, PyTorch 2.9.1+cu128, Triton 3.5.1. Earlier revisions of this file listed CUDA 12.4 / Python 3.11 / PyTorch 2.6 / Triton 3.2, which is the stack the superseded results were produced on.
• The head-to-head additionally needs pip install flash-maxsim late-interaction-kernels.

pip install -r requirements.txt

Paths are resolved from each script's own location, so any checkout works from any working directory. Data and output locations are overridable:

variable	meaning	default
TILEMAXSIM_DATA	datasets and prebuilt baseline indexes	<repo>/data
TILEMAXSIM_RESULTS	where re-run scripts write	<repo>/results
CONDA_ENV	optional; if set, baseline scripts activate it	unset

Reproducing the paper

# Every paper table, from one build. One process per section, see the warning below.
python src/regenerate_all_tables.py --only variants,baselines,dim_scaling,\
query_scaling,doc_scaling,tile_ablation,precision,roofline

# Head-to-head against the two published kernels. One process per shape.
python src/head2head_2026_09.py --list-shapes          # then one run per shape
python src/head2head_2026_09.py --one-shape 32,128,128,100000

# Embedding-dimension dispatch: correctness across d, and the speed of padding
# versus tiling at the same shape
python src/validate_dim_dispatch.py

# int64 validation, bandwidth sweep, dtype-matched baselines
python src/remeasure_2026_09.py --only all

Run one process per benchmark section and per shape. This is a correctness requirement, not a style preference. A third-party kernel that triggers cudaErrorIllegalAddress poisons the CUDA context, and every later measurement in that process then fails with the same error, including shapes that succeeded seconds earlier. In one run a torch.compile bug silently failed six whole table sections this way, and the failures looked like kernel defects but were not.

Measurement conventions, which the scripts enforce: FP16 inputs with FP32 accumulation; latency from torch.cuda.Event, median of 30 after 10 warmups; correctness checked against an exhaustive FP32 reference before any timing is believed; identical dtypes handed to every method, since a baseline that upcasts to FP32 is not a baseline. Counters come only from Nsight Compute, whose wall times are inflated roughly 1.6x by kernel replay and must never be read as latency.

Notes

• flash_maxsim_v2_multiquery in src/flash_maxsim_v2.py is the recommended entry point. It reads each document embedding from HBM once only when BLOCK_Nq >= N_q; the wrapper warns when it cannot, naming the number of passes. Heed the warning: it is the direct cause of the long-query losses above.
• Embedding dimensions at or below 128 are padded up to a power of two and the surplus lanes are masked, rather than routed to the dimension-tiled kernel. Masked lanes read as zero and contribute nothing to the dot product, so the two paths return bitwise identical scores; routing them to the tiled kernel instead halves the document tile and costs 1.75x at d=96 and 2.44x at d=100. Only d > 128 tiles.
• Pointer offsets are computed in int64. In int32 they overflowed at N_d * d >= 32,770, which every published shape avoided, but the largest by only 65,535 elements. Thanks to the contributor who reported it.
• The WARP and PLAID harnesses in baselines/ need local checkouts of the upstream repositories.

Citation

Please cite the paper rather than this repository. The arXiv v2 "Comments" field records which claims changed and why.