Fast Gauss Sums via Flash Attention

Nicolaj Rux; Sebastian Neumayer

cs.LGSep 4, 2026

Fast Gauss Sums via Flash Attention

Authors: Nicolaj Rux, Sebastian Neumayer

Organizations: Faculty of Mathematics Chemnitz University of Technology Reichenhainer Str. 39, 09126 Chemnitz, Germany

Abstract

Gaussian kernel sums are the computational core of maximum mean discrepancies (MMDs), kernel gradient flows, Stein variational gradient descent (SVGD), and many other kernel methods. At the same time, softmax attention has received an extraordinary amount of hardware-aware code engineering, culminating in flash attention. We show that Gauss kernel sums with arbitrary, signed weights can be evaluated via flash attention: two small input augmentations turn the normalized softmax reduction into the unnormalized Gauss sum, without writing a single line of custom GPU code. For feature dimension D>8 in fp16, this approach beats compiled PyTorch code as well as PyKeOps kernels (often significantly) in speed, memory-overhead and accuracy. Indeed, its memory scaling remains linear.

Figures & tables

#	Step	Work	Call
1	Input $k\in{\mathbb{R}}^{N\times D}$ , $q\in{\mathbb{R}}^{M\times D}$ , $v\in{\mathbb{R}}^{N\times C}$ , $\tau>0$	–
2	Output $s_{m}=\sum_{n=1}^{N}\Phi_{\tau}(q_{m},k_{n})v_{n}$ , $m=1,\dots,M$	–
3	$\tilde{v}_{n}\coloneqq v_{n}\exp(-\frac{\tau}{2}\\|k_{n}\\|^{2})$ , $n=1,\ldots,N$	$N(C+D)$	mult
4	Compute $\att_{\tau}(q,k,\tilde{v})\in{\mathbb{R}}^{M\times C}$ , $\lse_{\tau}(q,k)\in{\mathbb{R}}^{M}$	$MN(D{+}C)$	Flash
5	$s_{m}\coloneqq\att_{\tau}(q,k,\tilde{v})_{m}\exp(\lse_{\tau}(q,k)_{m})\exp(-\frac{\tau}{2}\\|q_{m}\\|^{2})$	$MC$	mult
6	return $s$

Algorithm 1 ( prescale ) Gaussian kernel summation via flash attention

#	Step	Work	Call
1	Input $k\in{\mathbb{R}}^{N\times D}$ , $q\in{\mathbb{R}}^{M\times D}$ , $v\in{\mathbb{R}}^{N\times C}$ , $\tau>0$ , $\kappa=\sqrt{N{+}1}$	–
2	Output $s_{m}=\sum_{n=1}^{N}\Phi_{\tau}(q_{m},k_{n})v_{n}$ , $m=1,\dots,M$	–
3	$\tilde{q}_{m}\coloneqq[q_{m},1,\frac{1}{2}\\|q_{m}\\|_{2}^{2}]\in{\mathbb{R}}^{D{+}2}$ , $m=1,\ldots,M$	$MD$	pad
4	$\tilde{k}_{n}\coloneqq[k_{n},-\frac{1}{2}\\|k_{n}\\|_{2}^{2},0],n=1,\ldots,N$ , $\tilde{k}_{0}\coloneqq e_{D+2}$	$ND$	pad
5	$\tilde{v}_{n}\coloneqq[v_{n},0],n=1,\ldots,N$ , $\tilde{v}_{0}\coloneqq\kappa\,e_{C+1}$	$NC$	pad
6	$[\alpha,\beta]\coloneqq\att_{\tau}(\tilde{q},\tilde{k},\tilde{v})$ with $\alpha\in{\mathbb{R}}^{M\times C}$ and $\beta\in{\mathbb{R}}^{M}$	$MN(D{+}C)$	Flash

Algorithm 2 ( reweight ) Gaussian kernel summation via flash attention

Panel of Figure 1 — Figure 1: Sweep over $N$ computing the forward of the Gauss sum ( 1 ) for fixed $B=4$ in fp16 with backend flash . The upper row is $D=3$ , $C=1$ and the lower row is $D=32$ , $C=32$ .

Panel of Figure 2 — Figure 2: Sweep over $N$ computing the forward and backward of the Gauss sum ( 1 ) for fixed $B=4$ , $D=32$ and $C=1$ in fp16 with backend flash .

Panel of Figure 3 — Figure 3: Sweep over $D$ computing the forward of the Gauss sum ( 1 ) for fixed $B=64$ , $C=1$ and $N=16384$ in fp16 with backend flash .

Panel of Figure 4 — Figure 4: Sweep over $D$ computing the forward of the Gauss sum ( 1 ) for fixed $B=64$ , $C=1$ and $N=16384$ in fp32 with backend memory_efficient .

Panel of Figure 5 — Figure 5: Sweep over $N$ computing the forward of the Gauss sum ( 1 ) for fixed $B=4$ , $D=32$ and $C=32$ in fp32 with backend memory_efficient .

Kernel Methods

Kernel methods use similarity functions and reproducing kernel spaces to model nonlinear relationships, quantify uncertainty, and compare data.

See all 227 papers

Flash-GMM: A Memory-Efficient Kernel for Scalable Soft Clustering
Jun 9, 2026Gal Bloch, Ariel Gera, Matan Orbach +2GPU Kernel Optimization GPU Acceleration
IBM Research
Projection-Free Transformers via Gaussian Kernel Attention
May 4, 2026Debarshi Kundu, Archisman Ghosh, Swaroop Ghosh +1Transformer Self-Attention
The Pennsylvania State University, University Park, USA
Broken Symmetry in BF16 Attention: Why FlashAttention Gradients Blow Up Late in Training
Sep 28, 2026Junlin Chen, Daize Dong, Huanwei Di +9Softmax Attention Efficient Attention
Rutgers University · Carnegie Mellon University · Oracle +2

Jun 9, 2026J/K move · Enter open · S save

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

Gal Bloch, Ariel Gera, Matan Orbach, Ohad Eytan+1

IBM Research

GPU Kernel Optimization GPU Acceleration High-Performance Computing Clustering Gaussian Mixture Models

Open paper AskPDF arXiv

We present 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.

Abstract

Kernel Methods

Explore similar work

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

Projection-Free Transformers via Gaussian Kernel Attention

Broken Symmetry in BF16 Attention: Why FlashAttention Gradients Blow Up Late in Training