FLASH-MAXSIM: IO-Aware Fused Kernels for Late-Interaction Scoring

Adi Raz Goldfarb; Daniel Ezer; Idan Friedman; Roi pony; Udi Barzelay

arxiv: 2605.29517 · v1 · pith:VZ5VDBWX · submitted 2026-05-28 · cs.IR

FLASH-MAXSIM: IO-Aware Fused Kernels for Late-Interaction Scoring

Roi Pony , Adi Raz Goldfarb , Idan Friedman , Daniel Ezer , Udi Barzelay This is my paper

Reviewed by Pith T0 review T1 audit T2 compute T3 formal T4 kernel 2026-06-29 05:36 UTCgrok-4.3pith:VZ5VDBWX record.json open to challenge →

classification cs.IR

keywords MaxSim operatorlate-interaction retrievalfused GPU kernelsIO-aware computationmemory optimizationColBERTtensor materialization

0 comments

The pith

Flash-MaxSim computes exact MaxSim scores without materializing the query-document similarity tensor by streaming tiles through on-chip SRAM and folding the reduction into the same pass.

A machine-rendered reading of the paper's core claim, the machinery that carries it, and where it could break.

The paper establishes that the MaxSim operator used in late-interaction models can produce identical scores without ever allocating the full similarity tensor, which otherwise reaches 21 GB for visual retrieval at scale. Query and document tiles move through SRAM while the per-row maximum is computed on the fly, and the same principle is extended to the backward pass, quantization, and padding-free inputs. A reader would care because the avoided tensor directly limits batch size and corpus scale on 40 GB GPUs in both inference and training, yet the method reports matching rankings with large measured gains in speed and memory use.

Core claim

Flash-MaxSim is an IO-aware fused GPU kernel that computes exactly the same MaxSim scores without ever materializing the tensor, by streaming query and document tiles through on-chip SRAM and folding the row-maximum reduction into the same pass; it extends the IO-aware principle through the training backward pass via an inverse-grid CSR construction that reuses the forward argmax for an atomic-free gradient reduction, and through INT8xINT8 quantization and variable-length scoring.

What carries the argument

The IO-aware fused GPU kernel that streams query and document tiles through on-chip SRAM while folding the row-maximum reduction into the same pass.

If this is right

Up to 3.9x faster inference on A100 and 4.7x on H100 than naive PyTorch at matched precision.
Up to 16x less inference memory and ~28x less training memory.
Unlocks corpus and batch sizes that exhaust standard PyTorch implementations.
Preserves 100% top-20 ranking agreement with an FP32 reference.
Supports INT8 quantization and padding-free variable-length scoring.

Where Pith is reading between the lines

These are editorial extensions of the paper, not claims the author makes directly.

The same tile-streaming and reduction-folding pattern could be applied to other operators that produce large intermediate matrices before a reduction step.
Memory savings of this magnitude may allow training or inference on hardware with smaller on-board memory without altering model architecture.
Avoiding materialization may reduce peak power draw during large-scale retrieval workloads even when raw speed is not the primary goal.

Load-bearing premise

The fused kernel reproduces the exact mathematical MaxSim operator, including correct argmax locations for the backward pass, with no numerical discrepancies and with tile sizes that fit in available SRAM for the target sequence lengths.

What would settle it

A side-by-side run on identical inputs that shows any difference in the final scores or in the argmax locations used for gradients between the fused kernel and a reference implementation that materializes the tensor.

Figures

Figures reproduced from arXiv: 2605.29517 by Adi Raz Goldfarb, Daniel Ezer, Idan Friedman, Roi pony, Udi Barzelay.

**Figure 1.** Figure 1: FLASH-MAXSIM streams Q and D tiles from HBM into SRAM, forms the sub-tile St and folds it into a running max on chip, and writes back only the scalar score. The [B, Lq, Ld] similarity tensor is never written to HBM. Algorithm 2 FLASH-MAXSIM forward (one program per (query, document) pair) 1: acc ← 0 ▷ FP32 register accumulator 2: for each query tile Qblk of Bq rows do 3: load Qblk into SRAM 4: m ← −∞ ▷ ru… view at source ↗

**Figure 2.** Figure 2: ColPali corpus scaling on A100-40GB (Lq=Ld=1024). Naive einsum in pure FP16 and in matched FP32 (FP32 accumulation) both hit OOM cliffs as the corpus grows (matched FP32 at B ≈7K, FP16 at B ≈20K); FLASH-MAXSIM keeps scaling to B ≈175K because its peak memory tracks only the document embeddings, so the Lq × Ld similarity matrix never exists. Left: latency. Right: peak memory [PITH_FULL_IMAGE:figures/full_f… view at source ↗

read the original abstract

Late-interaction retrieval (ColBERT, ColPali) scores a query against a document with the MaxSim operator: for every query token, the maximum similarity over the document tokens, summed over query tokens. The standard implementation materializes the full query-token x document-token similarity tensor in GPU memory; for visual ColPali at 10K documents this tensor alone is 21 GB in FP16, created only to be reduced to one score per document and discarded. It exhausts a 40 GB GPU and bounds the achievable batch size in both inference and training. We present Flash-MaxSim, an IO-aware fused GPU kernel that computes exactly the same scores without ever materializing the tensor, by streaming query and document tiles through on-chip SRAM and folding the row-maximum reduction into the same pass. We extend the IO-aware principle through the training backward pass, an inverse-grid CSR construction that reuses the forward argmax for an atomic-free, destination-owned gradient reduction, and through INT8xINT8 quantization and variable-length (padding-free) scoring. Flash-MaxSim is up to 3.9x faster on an A100 (4.7x on an H100) than naive PyTorch at matched precision, uses up to 16x less inference memory and ~28x less training memory, unlocks corpus and batch sizes that exhaust PyTorch entirely, preserves the exact ranking (100% top-20 agreement with an FP32 reference)

Editorial analysis

A structured set of objections, weighed in public.

Desk editor's note, referee report, simulated authors' rebuttal, and a circularity audit. Tearing a paper down is the easy half of reading it; the pith above is the substance, this is the friction.

Desk Editor's Note private letter to a colleague

Flash-MaxSim fuses MaxSim to skip materializing the similarity tensor via SRAM tiling and a CSR backward trick, delivering real memory and speed gains for ColBERT-style systems, but exact equivalence rests on ranking agreement rather than tight numerical checks.

read the letter

This paper gives a practical fused kernel for the MaxSim operator in late-interaction retrieval. It streams query and document tiles through on-chip SRAM, folds the row-max reduction into the same pass, and avoids ever building the full similarity tensor that blows up memory in standard PyTorch.

What is new is the extension of FlashAttention-style IO-aware fusion to this operator, including an inverse-grid CSR construction for the backward pass that reuses forward argmax locations for an atomic-free gradient reduction, plus support for INT8 quantization and padding-free variable lengths. The reported numbers are concrete: up to 3.9x faster inference on A100 (4.7x on H100), 16x less inference memory, and roughly 28x less training memory, with the ability to handle corpus and batch sizes that exhaust the naive approach.

The work does well on the engineering side. These are measurable wins for anyone running ColBERT or ColPali at scale in search or RAG, and the techniques address a real bottleneck without changing the mathematical operator.

The soft spot is verification. The abstract claims identical scores and 100% top-20 ranking agreement with an FP32 reference, but ranking agreement is a weak proxy. It does not rule out small per-score discrepancies or errors in argmax indices that could affect gradient flow during training. The stress-test concern about whether the tiled SRAM pass and CSR routing produce bit-identical results holds up here because the provided material gives no implementation details, benchmark methodology, or direct numerical comparisons beyond the ranking metric.

This paper is for engineers and practitioners who need to scale late-interaction models on commodity GPUs. A reader working on efficient retrieval implementations would get usable ideas and potential code to try.

It deserves peer review because the optimization is specific, the performance claims are testable, and the practical impact is clear even if more evidence on numerical fidelity would help.

Referee Report

2 major / 0 minor

Summary. The paper introduces Flash-MaxSim, an IO-aware fused GPU kernel for late-interaction retrieval (ColBERT, ColPali) that computes MaxSim scores by streaming query and document tiles through on-chip SRAM without materializing the full similarity tensor; it extends the approach to the training backward pass via inverse-grid CSR construction reusing forward argmax, to INT8 quantization, and to variable-length padding-free cases, claiming up to 3.9x (A100) / 4.7x (H100) speedup over naive PyTorch, 16x less inference memory, ~28x less training memory, and 100% top-20 ranking agreement with an FP32 reference.

Significance. If the exact numerical equivalence and gradient correctness hold, the work would enable substantially larger batch sizes and corpora in both inference and training for late-interaction models, particularly visual ones where the materialized tensor is prohibitive; the engineering focus on SRAM tiling and atomic-free backward reduction addresses a practical bottleneck in the field.

major comments (2)

[Abstract] Abstract: the central claim that the kernel 'computes exactly the same scores' and reuses 'identical argmax locations' for the backward pass rests solely on the reported 100% top-20 ranking agreement; no direct numerical comparison of per-document MaxSim scores, per-row maxima, or argmax indices against a materialized FP32 reference is provided, nor is any verification methodology or test for gradient correctness under the CSR construction described.
[Abstract] Abstract: the reported speedups and memory reductions are presented without benchmark methodology details (e.g., sequence lengths, batch sizes, hardware configuration, or comparison baselines beyond 'naive PyTorch'), making it impossible to assess whether the 3.9x/4.7x figures and 16x/28x memory claims are load-bearing or reproducible from the given information.

Simulated Author's Rebuttal

2 responses · 0 unresolved

We thank the referee for the careful reading and constructive comments. We address each major point below and will revise the manuscript accordingly to strengthen the claims with additional verification and clearer methodology.

read point-by-point responses

Referee: [Abstract] Abstract: the central claim that the kernel 'computes exactly the same scores' and reuses 'identical argmax locations' for the backward pass rests solely on the reported 100% top-20 ranking agreement; no direct numerical comparison of per-document MaxSim scores, per-row maxima, or argmax indices against a materialized FP32 reference is provided, nor is any verification methodology or test for gradient correctness under the CSR construction described.

Authors: We agree that direct numerical verification beyond ranking agreement would more rigorously support the exact-equivalence claim. The current manuscript validates via 100% top-20 agreement on retrieval tasks, which is the end-to-end metric of interest, but we will add an appendix with per-document MaxSim score comparisons (maximum absolute error tables), per-row argmax index matches, and a description of the gradient verification procedure (forward argmax reuse and CSR reduction checked against autograd on small tensors). revision: yes
Referee: [Abstract] Abstract: the reported speedups and memory reductions are presented without benchmark methodology details (e.g., sequence lengths, batch sizes, hardware configuration, or comparison baselines beyond 'naive PyTorch'), making it impossible to assess whether the 3.9x/4.7x figures and 16x/28x memory claims are load-bearing or reproducible from the given information.

Authors: The experimental section of the full manuscript specifies the configurations (query lengths 32-512, document lengths 128-2048, batch sizes 4-64, A100/H100 GPUs, FP16/INT8, and the exact PyTorch baseline implementation). However, these details are not summarized in the abstract. We will revise the abstract to include a brief parenthetical on the benchmark setup (e.g., "on sequences up to length 2048") and ensure the experimental section explicitly cross-references the reported numbers. revision: partial

Circularity Check

0 steps flagged

No circularity: engineering implementation with direct equivalence claim

full rationale

This is an engineering paper describing a fused GPU kernel implementation. The central claim is that the kernel computes exactly the same MaxSim scores as the materialized reference (verified by 100% top-20 ranking agreement). No mathematical derivation chain exists, no parameters are fitted and then called predictions, no self-citations are load-bearing for any uniqueness or ansatz, and no result reduces to its inputs by construction. The implementation is self-contained against external benchmarks (the standard MaxSim operator), satisfying the criteria for a non-circular finding.

Axiom & Free-Parameter Ledger

0 free parameters · 1 axioms · 0 invented entities

The paper contributes an optimized implementation of an existing operator rather than new mathematics or entities. No free parameters are introduced. The only background assumption is the standard definition of MaxSim.

axioms (1)

domain assumption MaxSim is defined as the sum over query tokens of the maximum similarity to any document token.
This is the established definition in ColBERT and ColPali models referenced in the abstract.

pith-pipeline@v0.9.1-grok · 5824 in / 1438 out tokens · 31091 ms · 2026-06-29T05:36:34.246338+00:00 · methodology

discussion (0)

Forward citations

Cited by 1 Pith paper

Reviewed papers in the Pith corpus that reference this work. Sorted by Pith novelty score.

TileMaxSim: IO-Aware GPU MaxSim Scoring with Dimension Tiling and Fused Product Quantization
cs.IR 2026-06 unverdicted novelty 7.0

TileMaxSim delivers IO-aware Triton kernels that reach 80% of peak HBM bandwidth for exact MaxSim scoring via multi-query tiling, dimension tiling, and fused product quantization, yielding 220x speedup over naive loop...

Reference graph

Works this paper leans on

2 extracted references · cited by 1 Pith paper

[1]

online" 'onlinestring :=

ENTRY address archivePrefix author booktitle chapter edition editor eid eprint eprinttype howpublished institution journal key month note number organization pages publisher school series title type volume year doi pubmed url lastchecked label extra.label sort.label short.list INTEGERS output.state before.all mid.sentence after.sentence after.block STRING...
[2]

write newline

" write newline "" before.all 'output.state := FUNCTION n.dashify 't := "" t empty not t #1 #1 substring "-" = t #1 #2 substring "--" = not "--" * t #2 global.max substring 't := t #1 #1 substring "-" = "-" * t #2 global.max substring 't := while if t #1 #1 substring * t #2 global.max substring 't := if while FUNCTION word.in bbl.in capitalize " " * FUNCT...

[1] [1]

online" 'onlinestring :=

ENTRY address archivePrefix author booktitle chapter edition editor eid eprint eprinttype howpublished institution journal key month note number organization pages publisher school series title type volume year doi pubmed url lastchecked label extra.label sort.label short.list INTEGERS output.state before.all mid.sentence after.sentence after.block STRING...

[2] [2]

write newline

" write newline "" before.all 'output.state := FUNCTION n.dashify 't := "" t empty not t #1 #1 substring "-" = t #1 #2 substring "--" = not "--" * t #2 global.max substring 't := t #1 #1 substring "-" = "-" * t #2 global.max substring 't := while if t #1 #1 substring * t #2 global.max substring 't := if while FUNCTION word.in bbl.in capitalize " " * FUNCT...