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

• URL: http://arxiv.org/abs/2605.29517v1
• Date: Thu, 28 May 2026 07:38:27 GMT
• Title: FLASH-MAXSIM: IO-Aware Fused Kernels for Late-Interaction Scoring
• Abstract summary: 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.
• Score: 2.159285655678094
• License: http://arxiv.org/licenses/nonexclusive-distrib/1.0/
• 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)

Related papers

• Nova: An End-to-End MLIR Compiler for Deep Learning [0.0]
  We show how Nova can fuse operations across operation boundaries, optimize complex memory hierarchies, and tuning execution down to the register level.<n>By reducing the memory footprint by up to 29% relative to PyTorch, Nova successfully trains a 144-million parameter model at 17,900 tokens/s.<n> Nova matches or modestly exceeds cuBLAS and XLA on TF32 matmuls on most shapes, maintaining a stringent 5e-4 relative error.
  arXiv  Detail & Related papers  (2026-07-15T05:16:28Z)
• 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)
• TileMaxSim: IO-Aware GPU MaxSim Scoring with Dimension Tiling and Fused Product Quantization [3.3723515662362265]
  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.
  arXiv  Detail & Related papers  (2026-06-24T23:03:56Z)
• MonaVec: A Training-Free Embedded Vector Search Kernel for Edge and Offline AI Systems [0.0]
  MonaVec is a deterministic, embedded vector-search kernel for edge and offline AI.<n>It targets on-device RAG, offline embedded retrieval -- the niche occupies for relational data: one file, one call, runs anywhere.
  arXiv  Detail & Related papers  (2026-06-17T18:00:54Z)
• Litespark Inference on Consumer CPUs: Custom SIMD Kernels for Ternary Neural Networks [0.00815557531820863]
  Litespark-Inference integrates directly with Hugging-Face, achieving 9.2x faster time-to-first-token, 52x higher throughput, and 14x memory reduction.<n>Our implementation, Litespark-Inference, is pip-installable and integrates directly with Hugging-Face, achieving 9.2x faster time-to-first-token, 52x higher throughput, and 14x memory reduction.
  arXiv  Detail & Related papers  (2026-05-07T16:07:39Z)
• ELSA: Exact Linear-Scan Attention for Fast and Memory-Light Vision Transformers [11.576817774473398]
  Existing attention accelerators often trade exact softmax semantics, depend on fused Core kernels, or incur sequential depth that limits FP32 throughput on long sequences.<n>We present bfELSA, an algorithmic reformulation of online softmax attention that reserves exact softmax semantics in real arithmetic with a emphprovable $mathcalO(ulog n)$ FP32 relative error bound.<n>Unlike FlashAttention-2/3, which rely on HMMA/GMMA Core instructions and provide no FP32 path, ELSA operates identically on A100s and resource-
  arXiv  Detail & Related papers  (2026-04-26T16:41:30Z)
• FlashSampling: Fast and Memory-Efficient Exact Sampling [62.5203057469482]
  FlashSampling is an exact sampling primitive that fuses sampling into the LM-head matmul and never materializes the logits tensor in.<n>Across H100, H200, B200, and B300 GPUs, FlashSampling speeds up kernel-level decode workloads.<n>In end-to-end vLLM experiments, it reduces time per output token by up to $19%$ on the models we test.
  arXiv  Detail & Related papers  (2026-03-16T19:37:08Z)
• FlashAttention-4: Algorithm and Kernel Pipelining Co-Design for Asymmetric Hardware Scaling [20.849774181776414]
  FlashAttention-4 achieves up to 1.3$times$ speedup over cuDNN 9.13 and 2.7$times$ over Triton on B200 GPUs with BF16.<n>We implement FlashAttention-4 entirely in CuTe- embedded in Python, achieving 20-30$times$ faster compile times compared to traditional C++ template-based approaches.
  arXiv  Detail & Related papers  (2026-03-05T18:24:49Z)
• 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)
• dParallel: Learnable Parallel Decoding for dLLMs [77.24184219948337]
  Diffusion large language models (dLLMs) offer parallel token prediction and lower inference latency.<n>Existing open-source models still require nearly token-length decoding steps to ensure performance.<n>We introduce dParallel, a simple and effective method that unlocks the inherent parallelism of dLLMs for fast sampling.
  arXiv  Detail & Related papers  (2025-09-30T16:32:52Z)
• Cut Your Losses in Large-Vocabulary Language Models [102.6981011879656]
  We propose Cut Cross-Entropy (CCE), a method that computes the cross-entropy loss without materializing the logits for all tokens into global memory.<n>CCE reduces the memory footprint of the loss from 24 GB to 1 MB, and the total training-time memory consumption of the head from 28 GB to 1 GB.
  arXiv  Detail & Related papers  (2024-11-13T20:30:15Z)
• Breaking the Memory Barrier: Near Infinite Batch Size Scaling for Contrastive Loss [59.835032408496545]
  We propose a tile-based strategy that partitions the contrastive loss calculation into arbitrary small blocks. We also introduce a multi-level tiling strategy to leverage the hierarchical structure of distributed systems. Compared to SOTA memory-efficient solutions, it achieves a two-order-of-magnitude reduction in memory while maintaining comparable speed.
  arXiv  Detail & Related papers  (2024-10-22T17:59:30Z)
• Hardware-Aware Parallel Prompt Decoding for Memory-Efficient Acceleration of LLM Inference [23.633481089469836]
  Auto-regressive decoding of Large Language Models (LLMs) results in significant overheads in their hardware performance.<n>We propose a novel parallel prompt decoding that requires only $0.0002$% trainable parameters, enabling efficient training on a single A100-40GB GPU in just 16 hours.<n>Our approach demonstrates up to 2.49$times$ speedup and maintains a minimal memory overhead of just $0.0004$%.
  arXiv  Detail & Related papers  (2024-05-28T22:19:30Z)
• HDCC: A Hyperdimensional Computing compiler for classification on embedded systems and high-performance computing [58.720142291102135]
  This work introduces the name compiler, the first open-source compiler that translates high-level descriptions of HDC classification methods into optimized C code. name is designed like a modern compiler, featuring an intuitive and descriptive input language, an intermediate representation (IR), and a retargetable backend. To substantiate these claims, we conducted experiments with HDCC on several of the most popular datasets in the HDC literature.
  arXiv  Detail & Related papers  (2023-04-24T19:16:03Z)

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.