The reviewed record of science sign in
Pith

arxiv: 2606.26439 · v1 · pith:GYPBEUWA · submitted 2026-06-24 · cs.IR · cs.DC· cs.PF

TileMaxSim: IO-Aware GPU MaxSim Scoring with Dimension Tiling and Fused Product Quantization

Ashutosh Sharma This is my paper

Reviewed by Pith T0 review T1 audit T2 compute T3 formal T4 kernel 2026-06-26 00:28 UTCgrok-4.3pith:GYPBEUWArecord.jsonopen to challenge →

classification cs.IR cs.DCcs.PF
keywords MaxSim scoringGPU kernel optimizationmulti-vector retrievalColBERTproduct quantizationIO-aware tilingHBM bandwidth utilizationTriton kernels
0
0 comments X

The pith

TileMaxSim kernels read each document embedding once from HBM and accumulate token maxima in registers to reach 80% peak bandwidth on H100 GPUs.

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

The paper shows that standard MaxSim implementations for models like ColBERT materialize an entire similarity matrix and therefore move the same document embeddings many times, using only 5-18% of available memory bandwidth. TileMaxSim replaces this with multi-query SRAM tiling that streams embeddings through shared memory exactly once while keeping per-query maxima in registers, plus dimension tiling for embeddings wider than 128 and fused product-quantization tables. On H100 hardware the kernels deliver 82 million documents per second and cut end-to-end scoring latency at 100 K candidates from 268 ms to 1.2 ms while producing identical rankings on MS MARCO and three BEIR sets.

Core claim

TileMaxSim is a family of Triton kernels that close the bandwidth gap in MaxSim scoring by (1) streaming document embeddings through shared memory while accumulating per-query-token maxima in registers, (2) partitioning the embedding dimension into 128-wide chunks, and (3) fusing product-quantization lookup tables in shared memory. These changes raise utilization to 80.2% of peak HBM bandwidth, yielding 82 M documents per second (71.6 M on real MS MARCO passages) and 220x speedup over loop-based code while preserving exact retrieval quality.

What carries the argument

Multi-query SRAM tiling that streams each document embedding from HBM exactly once while accumulating maxima in registers, augmented by 128-wide dimension tiling and shared-memory product-quantization tables.

If this is right

  • Scoring latency at 100 K candidates falls from 268 ms to 1.2 ms when TileMaxSim replaces the previous engine inside ColBERTv2/PLAID.
  • Throughput remains constant when the candidate set grows from 100 K to 500 K documents.
  • The same kernels support data-parallel sharding across multiple GPUs and run unchanged on embedding dimensions from 64 to 768 and in FP16, BF16, or FP32.
  • Exact ranking quality is preserved on MS MARCO and three BEIR benchmarks, so no re-ranking or accuracy loss occurs.
  • Fused product-quantization scoring reduces HBM traffic by up to 31x compared with materializing full embeddings.

Where Pith is reading between the lines

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

  • The same tiling pattern could be reused for other token-level similarity operations that currently materialize large intermediate matrices.
  • Because the kernel is written in Triton, the same approach may transfer to newer GPU architectures without hand-written CUDA.
  • Constant throughput up to 500 K candidates suggests the method could support larger candidate pools in retrieval pipelines that previously hit memory walls.

Load-bearing premise

The selected tile sizes and register accumulation strategy stay bandwidth-bound rather than compute-bound or register-pressure-bound for every query and document length that appears in practice.

What would settle it

A roofline plot or per-length timing table showing that throughput drops below the predicted bandwidth line for some common query or document lengths would falsify the central performance claim.

Figures

Figures reproduced from arXiv: 2606.26439 by Ashutosh Sharma.

Figure 1
Figure 1. Figure 1: Data flow of TileMaxSim V2-MQ. Query tiles are loaded once into SRAM; document tiles stream through SRAM one block at a time. Partial similarities are computed via Tensor Core tl.dot and reduced to running maxima in registers. No intermediate data is written back to HBM. the HBM IO of TileMaxSim V2-MQ with tile sizes BQ and BN is: IO =  𝑁𝑞 · 𝑑 |{z} read Q once + ⌈𝑁𝑞/BQ⌉ · 𝐵 · 𝑁𝑑 · 𝑑 | {z } read D per quer… view at source ↗
read the original abstract

Multi-vector retrieval models such as ColBERT achieve state-of-the-art accuracy through fine-grained token-level MaxSim scoring, yet existing GPU implementations leave most hardware performance unused. We give a roofline analysis of MaxSim on modern GPUs and identify a severe bandwidth gap: naive implementations reach only 5-18% of peak HBM bandwidth because they materialize the Nq x Nd similarity matrix, wasting memory traffic on data that is consumed once and discarded. We present TileMaxSim, a family of IO-aware Triton kernels that close this gap via (1) multi-query SRAM tiling that streams document embeddings through shared memory while accumulating per-query-token maxima in registers, reading each embedding from HBM exactly once; (2) dimension tiling that partitions the embedding dimension into 128-wide chunks, enabling scoring for d > 128 embeddings that overflow shared memory; and (3) fused product-quantization scoring via shared-memory lookup tables, cutting HBM I/O by up to ~31x. On NVIDIA H100 GPUs, TileMaxSim reaches 80.2% of peak HBM bandwidth and scores 82M documents/second (71.6M/s on real MS MARCO passages), a 220x speedup over loop-based scoring, 6.5x over fused PyTorch, 6.6-8.5x over torch.compile, and 469x the scoring throughput of WARP's CPU engine on the same node. TileMaxSim preserves exact retrieval quality: on MS MARCO and three BEIR benchmarks, rankings match reference MaxSim. As a drop-in replacement in ColBERTv2/PLAID, it cuts scoring latency at 100K candidates from 268 ms to 1.2 ms (98% lower end-to-end latency). We further show constant throughput from 100K to 500K documents, data-parallel multi-GPU sharding, robustness across dimensions 64-768, and FP16/BF16/FP32 support. Concurrent work independently develops an IO-aware fused MaxSim kernel; we differ in dimension tiling for d > 128 and fused product-quantization scoring.

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.

Referee Report

1 major / 2 minor

Summary. The paper presents TileMaxSim, a family of IO-aware Triton kernels for MaxSim scoring in multi-vector retrieval (e.g., ColBERT). It identifies a bandwidth gap in naive implementations that materialize the Nq×Nd similarity matrix, then proposes multi-query SRAM tiling with register accumulation (reading each embedding once from HBM), dimension tiling for d>128, and fused product-quantization via shared-memory tables. On H100 GPUs it reports 80.2% of peak HBM bandwidth, 82M documents/s (71.6M/s on MS MARCO), 220× speedup over loop-based scoring, and exact quality preservation on MS MARCO plus three BEIR sets; end-to-end latency at 100K candidates drops from 268 ms to 1.2 ms. Additional claims include constant throughput to 500K documents, multi-GPU scaling, and support for dimensions 64–768 and multiple precisions.

Significance. If the performance numbers hold under the stated assumptions, the work materially improves the deployability of state-of-the-art multi-vector retrievers by closing most of the gap between peak HBM bandwidth and realized MaxSim throughput. The concrete measured bandwidth utilization, throughput figures, and end-to-end latency reductions on standard benchmarks constitute a practical contribution; the dimension-tiling and fused-PQ extensions differentiate it from concurrent IO-aware kernels.

major comments (1)
  1. [Roofline analysis] Roofline analysis section: the central claim that the chosen tile sizes and register-accumulation strategy keep the kernel strictly bandwidth-bound (yielding the reported 80.2% HBM utilization) is not accompanied by per-query-length or per-document-length breakdowns, register-pressure measurements, or sensitivity sweeps over the length ranges actually present in MS MARCO (queries ~8–32 tokens, passages ~50–200 tokens) and BEIR. Because all headline throughput and speedup numbers are scaled from this single bandwidth figure, any length regime in which the working set exceeds registers or shared memory would render the 80.2% figure an optimistic point sample rather than a general property.
minor comments (2)
  1. [Results] The abstract and results sections report point estimates without error bars or multiple-run statistics; adding these would strengthen reproducibility of the 80.2% bandwidth and 82M docs/s figures.
  2. [Results] No per-query-length throughput tables or figures are provided despite the length-dependent nature of the tiling strategy; a small table or plot would clarify robustness.

Simulated Author's Rebuttal

1 responses · 0 unresolved

We thank the referee for the detailed and constructive review. The single major comment raises a valid point about the need for more granular validation of the bandwidth-bound claim. We address it below and will revise the manuscript accordingly.

read point-by-point responses

  1. Referee: [Roofline analysis] Roofline analysis section: the central claim that the chosen tile sizes and register-accumulation strategy keep the kernel strictly bandwidth-bound (yielding the reported 80.2% HBM utilization) is not accompanied by per-query-length or per-document-length breakdowns, register-pressure measurements, or sensitivity sweeps over the length ranges actually present in MS MARCO (queries ~8–32 tokens, passages ~50–200 tokens) and BEIR. Because all headline throughput and speedup numbers are scaled from this single bandwidth figure, any length regime in which the working set exceeds registers or shared memory would render the 80.2% figure an optimistic point sample rather than a general property.

    Authors: We agree that the current presentation would be strengthened by explicit length-specific breakdowns and sensitivity data. In the revised manuscript we will add a new figure and accompanying text that reports (1) measured HBM bandwidth utilization and achieved throughput for query lengths 8–32 and document lengths 50–200 tokens drawn from MS MARCO and the three BEIR sets, (2) register pressure and shared-memory occupancy for the tile sizes used in those regimes, and (3) a sweep confirming that the working set remains within the chosen register/shared-memory budgets across the observed length distribution. The headline 71.6 M docs/s figure is already measured on the actual MS MARCO passage distribution rather than a synthetic point sample; the additional plots will make this generality explicit rather than implicit. revision: yes

Circularity Check

0 steps flagged

No circularity; headline performance numbers are direct hardware measurements.

full rationale

The paper's central claims (80.2% of peak HBM bandwidth, 82M documents/second, 220x speedup) are presented as observed results from running TileMaxSim kernels on NVIDIA H100 GPUs, with explicit statements that rankings match reference MaxSim on MS MARCO and BEIR. No equations, fitted parameters, or self-citations are used to derive these quantities from the authors' own inputs. The roofline analysis is used only to identify a bandwidth gap in naive implementations; the reported utilization and speedups are not algebraically forced by that analysis or by any prior self-work. The implementation is self-contained against external benchmarks (real hardware counters and end-to-end latency on standard datasets).

Axiom & Free-Parameter Ledger

0 free parameters · 2 axioms · 0 invented entities

The central performance claims rest on standard GPU memory-hierarchy assumptions and the correctness of the Triton compiler rather than any new mathematical axioms or fitted constants.

axioms (2)
  • domain assumption Triton kernel compilation produces correct shared-memory and register tiling without hidden spills or synchronization errors
    Invoked implicitly when claiming 80.2% bandwidth utilization from the described tiling strategy.
  • domain assumption MaxSim semantics are preserved exactly when accumulation occurs in registers rather than materializing the full matrix
    Stated when claiming identical rankings on MS MARCO and BEIR.

pith-pipeline@v0.9.1-grok · 5932 in / 1450 out tokens · 22029 ms · 2026-06-26T00:28:45.647951+00:00 · methodology

discussion (0)

Sign in with ORCID, Apple, or X to comment. Anyone can read and Pith papers without signing in.

Reference graph

Works this paper leans on

29 extracted references · 5 canonical work pages · 3 internal anchors

  1. [1]

    ColBERT: Efficient and effective passage search via contextualized late interaction over BERT

    Omar Khattab and Matei Zaharia. ColBERT: Efficient and effective passage search via contextualized late interaction over BERT. InSIGIR, 2020

    2020

  2. [2]

    ColBERTv2: Effective and efficient retrieval via lightweight late interaction

    Keshav Santhanam, Omar Khattab, Jon Saad-Falcon, Christopher Potts, and Matei Zaharia. ColBERTv2: Effective and efficient retrieval via lightweight late interaction. InNAACL, 2022

    2022

  3. [3]

    PLAID: An efficient engine for late interaction retrieval

    Keshav Santhanam, Omar Khattab, Christopher Potts, and Matei Zaharia. PLAID: An efficient engine for late interaction retrieval. InCIKM, 2022

    2022

  4. [4]

    WARP: An efficient engine for multi-vector retrieval

    Jan Luca Scheerer, Matei Zaharia, Christopher Potts, Gustavo Alonso, and Omar Khattab. WARP: An efficient engine for multi-vector retrieval. InSIGIR, 2025

    2025

  5. [5]

    GEM: A native graph-based index for multi-vector retrieval

    Yao Tian, Zhoujin Tian, Xi Zhao, Ruiyuan Zhang, and Xiaofang Zhou. GEM: A native graph-based index for multi-vector retrieval. InSIGMOD, 2026

    2026

  6. [6]

    IGP: Efficient multi-vector retrieval via proximity graph index

    Zheng Bian, Man Lung Yiu, and Bo Tang. IGP: Efficient multi-vector retrieval via proximity graph index. InSIGIR, 2025

    2025

  7. [7]

    Fu, Stefano Ermon, Atri Rudra, and Christopher Ré

    Tri Dao, Daniel Y. Fu, Stefano Ermon, Atri Rudra, and Christopher Ré. FlashAtten- tion: Fast and memory-efficient exact attention with IO-awareness. InNeurIPS, 2022

    2022

  8. [8]

    FlashAttention-2: Faster attention with better parallelism and work partitioning

    Tri Dao. FlashAttention-2: Faster attention with better parallelism and work partitioning. InICLR, 2024

    2024

  9. [9]

    Flash-Decoding for long-context inference.Blog post, 2023

    Tri Dao, Daniel Haziza, Francisco Massa, and Grigory Sizov. Flash-Decoding for long-context inference.Blog post, 2023

    2023

  10. [10]

    FlashSinkhorn: IO-Aware Entropic Optimal Transport on GPU

    Felix X.-F. Ye, Xingjie Li, An Yu, Ming-Ching Chang, L. Chu, and D. Wertheimer. FlashSinkhorn: IO-aware entropic optimal transport on GPU.arXiv preprint arXiv:2602.03067, 2026

    work page internal anchor Pith review Pith/arXiv arXiv 2026

  11. [11]

    CAGRA: Highly parallel graph construction and approximate nearest neighbor search for GPUs.arXiv preprint arXiv:2308.15136, 2023

    Hiroyuki Ootomo, Akira Naruse, Corey Nolet, Ray Wang, Tamas Feher, and Yong Wang. CAGRA: Highly parallel graph construction and approximate nearest neighbor search for GPUs.arXiv preprint arXiv:2308.15136, 2023

    work page arXiv 2023

  12. [12]

    XTR: Rethinking the role of token retrieval in multi-vector retrieval

    Jinhyuk Lee, Zhuyun Dai, Sai Meher Karthik Duddu, Tao Lei, Iftekhar Naim, Ming-Wei Chang, and Vincent Zhao. XTR: Rethinking the role of token retrieval in multi-vector retrieval. InNeurIPS, 2023

    2023

  13. [13]

    MUVERA: Multi-vector retrieval via fixed dimensional encodings

    Laxman Dhulipala, Majid Hadian, Rajesh Jayaram, Jason Lee, and Vahab Mirrokni. MUVERA: Multi-vector retrieval via fixed dimensional encodings. InNeurIPS, 2024

    2024

  14. [14]

    Scaling ColPali to billions of PDFs with Vespa.Vespa Blog, 2024

    Jo Kristian Bergum. Scaling ColPali to billions of PDFs with Vespa.Vespa Blog, 2024

    2024

  15. [15]

    Storage access opti- mization for efficient GPU-centric information retrieval.The Journal of Super- computing, 81(4):613, 2025

    Susav Shrestha, Aayush Gautam, and Narasimha Reddy. Storage access opti- mization for efficient GPU-centric information retrieval.The Journal of Super- computing, 81(4):613, 2025

    2025

  16. [16]

    ColBERT-serve: Efficient multi-stage memory-mapped scoring

    Kaili Huang, Thejas Venkatesh, Uma Dingankar, Antonio Mallia, Daniel Campos, Jian Jiao, Christopher Potts, Matei Zaharia, Kwabena Boahen, Omar Khattab, Saarthak Sarup, and Keshav Santhanam. ColBERT-serve: Efficient multi-stage memory-mapped scoring. InECIR, 2025

    2025

  17. [17]

    ESPN: Memory-efficient multi-vector information retrieval

    Susav Shrestha, Narasimha Reddy, and Zongwang Li. ESPN: Memory-efficient multi-vector information retrieval. InISMM, 2024

    2024

  18. [18]

    Sculpting the Vector Space: Towards Efficient Multi-Vector Visual Document Retrieval via Prune-then-Merge Framework

    Yibo Yan, Mingdong Ou, Yi Cao, Xin Zou, Jiahao Huo, Shuliang Liu, James Kwok, and Xuming Hu. Sculpting the vector space: Towards efficient multi-vector visual document retrieval via prune-then-merge framework.arXiv preprint arXiv:2602.19549, 2026

    work page internal anchor Pith review Pith/arXiv arXiv 2026

  19. [19]

    Scaling attention via feature sparsity

    Yan Xie, Tiansheng Wen, Tangda Huang, Bo Chen, Chenyu You, Stefanie Jegelka, and Yifei Wang. Scaling attention via feature sparsity. InICLR, 2026

    2026

  20. [20]

    AutoKernel: Autonomous GPU kernel optimization via iterative agent-driven search.arXiv preprint arXiv:2603.21331, 2026

    Jaber Jaber and Osama Jaber. AutoKernel: Autonomous GPU kernel optimization via iterative agent-driven search.arXiv preprint arXiv:2603.21331, 2026

    work page arXiv 2026

  21. [21]

    Billion-scale similarity search with GPUs.IEEE Transactions on Big Data, 7(3):535–547, 2021

    Jeff Johnson, Matthijs Douze, and Hervé Jégou. Billion-scale similarity search with GPUs.IEEE Transactions on Big Data, 7(3):535–547, 2021

    2021

  22. [22]

    Sparton: Fast and memory-efficient Triton kernel for learned sparse retrieval

    Thong Nguyen, Cosimo Rulli, Franco Maria Nardini, Rossano Venturini, and Andrew Yates. Sparton: Fast and memory-efficient Triton kernel for learned sparse retrieval. InSIGIR, 2026

    2026

  23. [23]

    Meta Andromeda: Supercharging Advantage+ automation with the next-gen personalized ads retrieval engine.Meta Engineering Blog, 2024

    Chonglin Sun, Nancy Yu, Haiyu Lu, Liang Wang, Yunchen Pu, Gaoxiang Liu, Neeraj Bhatia, and GP Musumeci. Meta Andromeda: Supercharging Advantage+ automation with the next-gen personalized ads retrieval engine.Meta Engineering Blog, 2024

    2024

  24. [24]

    VecFlow-Chamfer: A GPU-based data management system for high-performance multi-vector search on superchips

    Chenghao Mo, Ben Karsin, Philip Adams, and Minjia Zhang. VecFlow-Chamfer: A GPU-based data management system for high-performance multi-vector search on superchips. InSIGMOD, 2026

    2026

  25. [25]

    GIGP+: A CPU-GPU co-processing engine for multi-vector retrieval

    Zheng Bian, Man Lung Yiu, and Bo Tang. GIGP+: A CPU-GPU co-processing engine for multi-vector retrieval. InSIGIR, 2026

    2026

  26. [26]

    Fu, Hermann Kumbong, Eric Nguyen, and Christopher Ré

    Daniel Y. Fu, Hermann Kumbong, Eric Nguyen, and Christopher Ré. FlashFFT- Conv: Efficient convolutions for long sequences with Tensor Cores. InICLR, 2024

    2024

  27. [27]

    Efficient multi- vector dense retrieval with bit vectors

    Franco Maria Nardini, Cosimo Rulli, and Rossano Venturini. Efficient multi- vector dense retrieval with bit vectors. InECIR, 2024

    2024

  28. [28]

    Fu, Tri Dao, Khaled K

    Daniel Y. Fu, Tri Dao, Khaled K. Saab, Armin W. Thomas, Atri Rudra, and Christo- pher Ré. Hungry Hungry Hippos: Towards language modeling with state space models. InICLR, 2023

    2023

  29. [29]

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

    Roi Pony, Daniel Ezer, Adi Raz Goldfarb, Idan Friedman, Oshri Naparstek, and Udi Barzelay. Flash-MaxSim: IO-aware fused kernels for late-interaction retrieval. arXiv preprint arXiv:2605.29517, 2026

    work page internal anchor Pith review Pith/arXiv arXiv 2026