Flash-GMM: A Memory-Efficient Kernel for Scalable Soft Clustering

• URL: http://arxiv.org/abs/2606.10896v1
• Date: Tue, 09 Jun 2026 14:07:18 GMT
• Title: Flash-GMM: A Memory-Efficient Kernel for Scalable Soft Clustering
• Abstract summary: Flash-GMM achieves a textbf20$times$ speedup over existing implementations.<n>We show that soft GMM clustering is now a viable drop-in replacement for $k$-means.
• Score: 3.845089876398699
• License: http://creativecommons.org/licenses/by-nc-nd/4.0/
• Abstract: We present \textbf{Flash-GMM}, a fused Triton kernel for efficient computation of Gaussian Mixture Models (GMMs) over large-scale data in a single GPU pass. By eliminating the need to materialize the full responsibility matrix in GPU memory, Flash-GMM achieves a \textbf{20$\times$} speedup over existing implementations and enables training on datasets more than \textbf{100$\times$} larger than previously feasible on one device. To demonstrate its impact, we integrate Flash-GMM into the IVF coarse quantizer for approximate nearest-neighbor (ANN) search. We show that soft GMM clustering is now a viable drop-in replacement for $k$-means, and that GMM responsibilities can be leveraged to assign border vectors to multiple clusters. Our approach reaches fixed recall targets with up to $1.7\times$ fewer distance computations, or equivalently, yields $+2$--$12$ recall@10 at matched computational cost. We release the kernel as an open-source project.

Related papers

• FluxBin: Flexible LUT-based Ultra-low-bit LLM Inference by Algorithm-Kernel Synergy [32.21687694159134]
  FluxBin is an algorithm that synerizes binary quantization with a highly optimized kernel.<n>FluxBin achieves up to $5.92times speedup and $10.19times energy savings.<n>This effectively enables the deployment of 70B-scale models on one single A100 bases with a $4times memory reduction.
  arXiv  Detail & Related papers  (2026-08-16T08:01:06Z)
• Accelerating GPU Inference of Large Language Models with Moderately Unstructured Sparse Weight Matrices [18.428868416628017]
  This paper proposes an efficient inference method for large language models (LLMs) with moderate sparsity.<n>It achieves up to 1.64x-level speedup over SpInfer and up to 1.41x end-to-end speedups over FlashLLM.
  arXiv  Detail & Related papers  (2026-06-13T13:38:27Z)
• On Efficient Scaling of GNNs via IO-Aware Layers Implementations [1.0262304700896199]
  Graph Neural Networks (GNNs) are bottlenecked by irregular memory access.<n>We show that widely used layers fall into three kernel families: SpMM-based convolutions, reduction-based aggregations, and attention-based layers.<n>For each family, we develop kernels that reduce data movement, improve locality, and remain robust across realistic graphs.
  arXiv  Detail & Related papers  (2026-05-29T16:22:45Z)
• Stochastic Sparse Attention for Memory-Bound Inference [19.301894658575502]
  We present Additive No-mulT Attention (SANTA), a method that sparsifies value-cache access by sampling $S ll n_k$ indices from the post-softmax distribution.<n>We also propose Bernoulli $qKmathsfT$ sampling as a complementary technique to sparsify the score stage.
  arXiv  Detail & Related papers  (2026-05-03T14:44:14Z)
• 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)
• SonicMoE: Accelerating MoE with IO and Tile-aware Optimizations [54.303301888915406]
  Mixture of Experts (MoE) models have emerged as the de facto architecture for scaling up language models without significantly increasing the computational cost.<n>We propose a memory-efficient algorithm to compute the forward and backward passes of MoEs with minimal activation caching.<n>We also propose a novel "token rounding" method that minimizes the wasted compute due to padding in Grouped GEMM kernels.
  arXiv  Detail & Related papers  (2025-12-16T04:39:10Z)
• FlashMoE: Fast Distributed MoE in a Single Kernel [1.866526462692252]
  FlashMoE is a fully GPU-resident MoE operator that fuses expert computation and inter-GPU communication into a single persistent GPU kernel.<n>We show that FlashMoE achieves up to 9x higher GPU utilization, 6x lower latency, 5.7x higher throughput, and 4x better overlap efficiency compared to state-of-the-art baselines.
  arXiv  Detail & Related papers  (2025-06-05T06:29:14Z)
• Sublinear Variational Optimization of Gaussian Mixture Models with Millions to Billions of Parameters [5.429282997550318]
  We train GMMs with over 10 billion parameters on about 100 million images, and observe training times of approximately nine hours on a single state-of-the-art CPU.<n>Our proposed algorithm significantly reduces runtime complexity per iteration from $mathcalO(NCD2)$ to a complexity scaling linearly with $D$ and remaining constant w.r.t.<n>As a proof of concept, we train GMMs with over 10 billion parameters on about 100 million images, and observe training times of approximately nine hours on a single state-of-the-art CPU.
  arXiv  Detail & Related papers  (2025-01-21T17:11:25Z)
• 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)
• 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 Dataset Distillation Using Random Feature Approximation [109.07737733329019]
  We propose a novel algorithm that uses a random feature approximation (RFA) of the Neural Network Gaussian Process (NNGP) kernel. Our algorithm provides at least a 100-fold speedup over KIP and can run on a single GPU. Our new method, termed an RFA Distillation (RFAD), performs competitively with KIP and other dataset condensation algorithms in accuracy over a range of large-scale datasets.
  arXiv  Detail & Related papers  (2022-10-21T15:56:13Z)
• Giga-scale Kernel Matrix Vector Multiplication on GPU [19.663081364196778]
  Kernel matrix-vector multiplication (KMVM) is a foundational operation in machine learning and scientific computing.<n>As KMVM tends to scale quadratically in both memory and time, applications are often limited by these computational constraints.<n>We propose a novel approximation procedure coined textitFaster-Fast and Free Memory Method ($fthreem$) to address these scaling issues.
  arXiv  Detail & Related papers  (2022-02-02T15:28:15Z)
• Coded Stochastic ADMM for Decentralized Consensus Optimization with Edge Computing [113.52575069030192]
  Big data, including applications with high security requirements, are often collected and stored on multiple heterogeneous devices, such as mobile devices, drones and vehicles. Due to the limitations of communication costs and security requirements, it is of paramount importance to extract information in a decentralized manner instead of aggregating data to a fusion center. We consider the problem of learning model parameters in a multi-agent system with data locally processed via distributed edge nodes. A class of mini-batch alternating direction method of multipliers (ADMM) algorithms is explored to develop the distributed learning model.
  arXiv  Detail & Related papers  (2020-10-02T10:41:59Z)

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.