# Fast Gauss Sums via Flash Attention

arXiv: https://arxiv.org/abs/2609.04910 · HTML: https://mlsanity.com/paper/2609.04910/

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Builds on 8 of its 13 references that are in the map.
- First posted on arXiv 4 Sep 2026.

## Builds on (most cited first, 8 of 8)

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness](https://mlsanity.com/paper/2205.14135/) (2022)
- [FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioning](https://mlsanity.com/paper/2307.08691/) (2023)
- [Stein Variational Gradient Descent: A General Purpose Bayesian Inference Algorithm](https://mlsanity.com/paper/1608.04471/) (2016)
- [Maximum Mean Discrepancy Gradient Flow](https://mlsanity.com/paper/1906.04370/) (2019)
- [Kernel Operations on the GPU, with Autodiff, without Memory Overflows](https://mlsanity.com/paper/2004.11127/) (2020)
- [Generative Sliced MMD Flows with Riesz Kernels](https://mlsanity.com/paper/2305.11463/) (2023)
- [Fast Kernel Summation in High Dimensions via Slicing and Fourier Transforms](https://mlsanity.com/paper/2401.08260/) (2024)

## Abstract (arXiv)

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.