FlashDiffusion: Fused Tiled Kernel Spectral Decomposition

Julio Candanedo

cs.LGSep 18, 2026

FlashDiffusion: Fused Tiled Kernel Spectral Decomposition

Authors: Julio Candanedo

Organizations: SparseTrace LLC Appleton, WI 54913, USA

Abstract

Diffusion maps, and kernel methods more generally, provide an interpretable nonlinear spectral representation basis for geometric learning. In the geometric limit, small bandwidth, these matrices tend to be high rank and thus require materializing dense Gaussian kernels requires $O(N^2)$ memory. We introduce FlashDiffusion, a matrix-free method that evaluates dense Gaussian kernel blocks in fused GPU tiles and couples the eigensolver to an empirical $β$ -flow that selects the finite-sample resolution scale. A continuation over sample size and bandwidth warm-starts increasingly expensive spectral solves from coarser resolutions.

Figures & tables

Panel of Figure 1 — Figure 1: FlashDiffusion on uniform samples of the flat torus $T^{6}$ (RTX PRO 6000). (a) Wall time versus sample count; the dashed guide is exact $\mathcal{O}(N^{2})$ scaling, and the flat region below $N\approx 2^{14}$ is the fixed launch/solver overhead. Four-GPU points continue the curve to $N=2^{22}\approx 4.2$ M. (b) The bandwidth emitted by the $(n,\beta)$ ladder, against the Singer exponent $2/(d+6)=1/6$ for $d=6$ (dashed, anchored at the first point). (c) Recovery of the first harmonic band: projection error of the leading nontrivial eigenspace onto the exact $T^{6}$ harmonics, decaying with sample count toward the resolution floor. Error bars are one standard deviation over repeated runs where multiple runs were performed.

Panel of Figure 2 — Figure 2: Ramachandran free-energy distributions for alanine dipeptide. (a) Reference distribution, (b) reconstruction using 8 diffusion-map (DMAP) modes, and (c) reconstruction using 256 DMAP modes. The horizontal coordinate is the dihedral angle $\psi$ and the vertical coordinate is $\phi$ ; both span $[-180^{\circ},180^{\circ}]$ , with $(\phi,\psi)=(0,0)$ at the center of each panel. The shared color bar reports the free energy $F/k_{B}T$ from 0 to 8.

Figure 3: Spectral resolution of alanine-dipeptide coordinates. (a) Mean absolute wrapped error for the backbone dihedrals $\phi$ and $\psi$ as the number of retained diffusion modes $m$ increases. Solid curves evaluate the completely held-out trajectory, while dashed curves evaluate $5{,}000$ randomly held-out frames from the decoder-training trajectories. The vertical markers indicate the ambient feature dimension $D=80$ and the detected onset $m_{\psi}=16$ of the principal $\psi$ -resolution regime. (b) Held-out Ramachandran free-energy surface at $\beta^{\star}=276.4$ , with random-holdout frames colored by joint angular reconstruction error and a strided segment of the decoded held-out trajectory overlaid as a curve.

Appendix figures & tables3 assets

Supplementary material from the paper’s appendix.

Appendix

$k$	retained $f_{k}$	$1-e^{-c_{2}k}$
32	2.2%	2.18%
64	4.4%	4.31%
128	8.6%	8.43%
256	16.4%	16.2%
1024	50.4%	50.7%

Appendix

Table 1: Gaussian kernel row mass retained under fixed-

k

truncation on a unit-variance Swiss roll with

N=2\times 10^{5}

and

\beta^{\star}=8.75

, averaged over 400 rows. The local

d=2

prediction

1-e^{-c_{2}k}

uses

c_{2}\approx 6.9\times 10^{-4}

$k$	$\lambda_{\min}(\mathsf{K}\odot A_{k})$	neg. spectral mass	$\lambda_{\min}(A_{k})$
16	$-4.10$	26.9%	$-4.79$
32	$-5.84$	27.8%	$-7.85$
64	$-7.71$	25.8%	$-13.44$
128	$-7.58$	19.0%	$-33.07$
256	$-2.73$	9.0%	$-61.03$

Appendix

Table 2: Spectrum of a symmetrically

k

NN-truncated Gaussian kernel on the Swiss roll with

N=3000

\beta^{\star}=2.158

, and neighborhood masks symmetrized by union. Negative spectral mass is

\sum_{\lambda<0}|\lambda|/\sum|\lambda|

. The final column reports the smallest eigenvalue of the mask itself.

Panel of Figure 4 — Figure 4: Raw matrix-free Gaussian-kernel scaling on the centered anisotropic flat torus $T^{6}$ ( $D=32$ , $\beta=1$ ) using one RTX–PRO–6000 GPU. (a) True matrix–vector application $Kv$ ( $b=1$ ). (b) Eight-right-hand-side matrix–matrix application $KV$ ( $b=8$ ). XGEMM –TILT and KeOps evaluate the same dense Gaussian interaction without materializing the $N\times N$ kernel matrix. Points show mean wall time and error bars show one standard deviation over repeated timed applications; the dashed guides indicate exact $\mathcal{O}(N^{2})$ scaling in the large- $N$ regime.

Spectral Methods

Spectral methods use eigenvalues, eigenvectors, and frequency decompositions to analyze data, graphs, operators, and learning algorithms.

See all 235 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
GeoShrink: Accelerating Diffusion Transformers with Two Lines of Code
Sep 27, 2026Haosen Li, Wenshuo Chen, Shaofeng Liang +3Video Diffusion Models Diffusion Transformer
The Hong Kong University of Science and Technology (Guangzhou) · Griffith University · Data61, CSIRO +1
Fast Gauss Sums via Flash Attention
Sep 4, 2026Nicolaj Rux, Sebastian NeumayerKernel Methods
Faculty of Mathematics Chemnitz University of Technology Reichenhainer Str. 39, 09126 Chemnitz, Germany

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

Spectral Methods

Explore similar work

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

GeoShrink: Accelerating Diffusion Transformers with Two Lines of Code

Fast Gauss Sums via Flash Attention