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

• URL: http://arxiv.org/abs/2606.26439v1
• Date: Wed, 24 Jun 2026 23:03:56 GMT
• Title: TileMaxSim: IO-Aware GPU MaxSim Scoring with Dimension Tiling and Fused Product Quantization
• Abstract summary: Multi-vector retrieval models such as ColBERT achieve state-of-the-art accuracy through fine-grained token-level MaxSim scoring.<n>Existing GPU implementations leave most hardware performance unused.<n>We present TileMaxSim, a family of IO-aware Triton kernels that close this gap.
• Score: 3.3723515662362265
• License: http://arxiv.org/licenses/nonexclusive-distrib/1.0/
• 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.

Related papers

• SparseDitto: Customizing GPU Kernels for Different Sparsity Patterns with LLM-Based Agentic System [45.52961410486319]
  We present SparseDitto, an LLM-based system that constructs a GPU kernel for each matrix, operator, and target GPU.<n>SparseDitto supports SpMV, SpMM, and SpGEMM within a unified design framework.<n>Its generated SpMM kernels also accelerate full-batch GCN training by up to 3.39x.
  arXiv  Detail & Related papers  (2026-08-05T16:41:03Z)
• A Modern Multimodal Assistant on a 6 GB 2011 GPU: Stage-Validated, All-GPU CUDA Inference for Fermi [0.0]
  This report keeps the hardware and asks what a model that fits can do.<n>We deploy Mini-V-4.6, a modern multimodal assistant pairing a SigLIP2 encoder and window-attention merger.<n>The system answers an image question end-to-end in 1.7s.
  arXiv  Detail & Related papers  (2026-07-16T04:48:44Z)
• Mixture-of-Parallelisms: Towards Memory-Efficient Training Stack for Mixture-of-Experts Models [67.17268530365189]
  This paper showcases a memory-efficient training stack for Mixture-of-Experts (MoE) models.<n>It combines and specializes various existing and novel parallelism techniques at different layers and stages of the MoE model training pipeline.
  arXiv  Detail & Related papers  (2026-07-02T08:06:57Z)
• GPUSparse: GPU-Accelerated Learned Sparse Retrieval with Parallel Inverted Indices [3.3723515662362265]
  We present GPUSparse, a system for GPU-accelerated exact learned sparse retrieval.<n>A GPU-parallel inverted index with block-aligned, warp-coalesced posting lists is used.<n>Batched scatter-add scoring algorithm processes hundreds of queries simultaneously.
  arXiv  Detail & Related papers  (2026-06-24T23:07:57Z)
• FLASH-MAXSIM: IO-Aware Fused Kernels for Late-Interaction Scoring [2.159285655678094]
  Flash-MaxSim is an IO-aware kernel that computes exactly the same scores without ever materializing the tensor.<n>It is up to 3.9x faster on an A100 than naive PyTorch at matched precision, uses up to 16x less inference memory and 28x less training memory.
  arXiv  Detail & Related papers  (2026-05-28T07:38:27Z)
• A Controlled Study of Memory Hierarchy Transitions in Quantum Circuit Simulation on Apple M4 Pro Unified Memory Architecture [0.0]
  State-vector quantum circuit simulation is memory-bandwidth bound.<n>We address this using the Apple M4 Pro Unified Memory Architecture.<n>We show that peak streaming bandwidth does not predict simulation speedup for non-contiguous memory access patterns.
  arXiv  Detail & Related papers  (2026-05-09T08:22:55Z)
• ZenBrain: A Neuroscience-Inspired 7-Layer Memory Architecture for Autonomous AI Systems [51.56484100374058]
  On LongMemEval-500, ZenBrain matches a long-context oracle's binary-judge accuracy to within 4.5 pp.<n> ZenBrain is a 7-layer neuroscience-inspired memory architecture.
  arXiv  Detail & Related papers  (2026-04-26T20:39:19Z)
• Scaling DoRA: High-Rank Adaptation via Factored Norms and Fused Kernels [83.99688944263843]
  Weight-De Low-Rank Adaptation (DoRA) extends LoRA by decoupling weight magnitude from direction.<n>At d_in = 8192 and rank r = 384, a single module's norm requires about 512 MB of transient working memory in bf16.<n>A factored norm decomposes the squared norm into base, cross, and Gram terms computable through O(d_out r + r2) intermediates, eliminating the dense product.
  arXiv  Detail & Related papers  (2026-03-23T17:57:24Z)
• Chronicals: A High-Performance Framework for LLM Fine-Tuning with 3.51x Speedup over Unsloth [0.0]
  We present Chronicals, an open-source training framework achieving 3.5x speedup over Unsloth.<n>We provide complete mathematical foundations: online softmax correctness, FlashAttention IO complexity O(N2 d2 M-1), LoRA+ learning rate gradient approximations.
  arXiv  Detail & Related papers  (2026-01-06T00:00:55Z)
• FuseSampleAgg: Fused Neighbor Sampling and Aggregation for Mini-batch GNNs [51.56484100374058]
  FuseSampleAgg fuses neighbor and sampling mean aggregation into a single pass for GraphSAGE.<n>Operator is deterministic, integrates with standard PyTorchs, and ships with scripts that reproduce all tables and figures from CSV logs.
  arXiv  Detail & Related papers  (2025-11-17T17:57:18Z)
• 70% Size, 100% Accuracy: Lossless LLM Compression for Efficient GPU Inference via Dynamic-Length Float [52.079202872069835]
  Large-scale AI models, such as Large Language Models (LLMs) and Diffusion Models (DMs) have grown rapidly in size.<n>We introduce Dynamic-Length Float (DFloat11), a compression framework that reduces LLM and DM size by 30% while preserving outputs that are bit-for-bit identical to the original model.
  arXiv  Detail & Related papers  (2025-04-15T22:38:38Z)
• vTensor: Flexible Virtual Tensor Management for Efficient LLM Serving [53.972175896814505]
  Large Language Models (LLMs) are widely used across various domains, processing millions of daily requests. Large Language Models (LLMs) are widely used across various domains, processing millions of daily requests.
  arXiv  Detail & Related papers  (2024-07-22T14:37:58Z)
• Efficient and Generic 1D Dilated Convolution Layer for Deep Learning [52.899995651639436]
  We introduce our efficient implementation of a generic 1D convolution layer covering a wide range of parameters. It is optimized for x86 CPU architectures, in particular, for architectures containing Intel AVX-512 and AVX-512 BFloat16 instructions. We demonstrate the performance of our optimized 1D convolution layer by utilizing it in the end-to-end neural network training with real genomics datasets.
  arXiv  Detail & Related papers  (2021-04-16T09:54:30Z)

This list is automatically generated from the titles and abstracts of the papers in this site.

This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.