Flash-GMM

KO | EN

gitlite — search

#python #agentic-ai #ai-agent #automation #genpark #llm-tools #mcp-server #dsh-plugin #deepseek-harness #dsh #llm #ai-agents

★ 24 Open GitHub ↗

Flash-GMM, a memory-efficient GPU kernel that unlocks GMM soft clustering at billion-scale

// repository documentation

Was this content helpful?

★ 0 (0 ratings)

# Flash-GMM

A memory-efficient, IO-aware Triton GPU kernel for the Gaussian Mixture Model (GMM) E-step at arbitrary scale ([Bloch et al. 2026](https://arxiv.org/abs/2606.10896)).

## Overview

Computing GMM responsibilities for large datasets is memory-bound: a naive implementation materializes an N×K responsibility matrix that exhausts GPU memory well before the scales where modern applications operate. Flash-GMM eliminates this matrix entirely.

Inspired by the IO-aware tiling strategy of [FlashAttention](https://arxiv.org/abs/2205.14135), Flash-GMM computes the full GMM E-step in a **single pass** over the data, accumulating only O(KD) sufficient statistics. The N×K responsibility matrix is never written to memory.

## Key Properties

| Property | Value |
|---|---|
| Kernel memory (N=1M, K=1024, D=128) | **4.5 MB** |
| TorchGMM memory (same config) | 21,006 MB |
| Memory reduction | **4,668×** |
| Speedup vs SciPy (CPU) | **766–1,740×** |
| Speedup vs TorchGMM (GPU) | **19–32×** |
| Max N on A100 80GB (streaming) | **1B+** |
| Validated on | A100 80GB, H100, RTX 5080 |

## How It Works

The kernel processes data in tiles of `BLOCK_N` rows. For each tile:

1. **Pass 1 (log-sum-exp)**: Loads the tile into registers and computes the per-sample log-normaliser `log Z_i` via a numerically stable online log-sum-exp, iterating over all K components in tiles of `BLOCK_K`.

2. **Pass 2 (accumulation)**: Reuses the tile from registers (no second HBM read), computes responsibilities `r_ik = exp(log z_ik - log Z_i)`, and atomically accumulates the sufficient statistics `N_k`, `mu_acc`, `sig_acc` into O(KD) global buffers.

The tile never leaves on-chip memory between the two passes, giving a **single HBM read of X** per iteration. The N×K responsibility matrix is discarded immediately after each tile.

## Installation

```bash
pip install torch triton
```

No other dependencies required for the kernel itself.

## Usage

```python
import math
import torch
from flash_gmm import flash_gmm_estep

# Inputs (all GPU tensors, float32)
N, K, D = 1_000_000, 1024, 128
X            = torch.randn(N, D, device='cuda')        # data
mu           = torch.randn(K, D, device='cuda')        # component means
log_sigma_sq = torch.zeros(K, device='cuda')           # log variances
log_pi       = torch.full((K,), -math.log(K), device='cuda')  # log weights

# E-step
logZ, Nk, mu_acc, sig_acc = flash_gmm_estep(X, mu, log_sigma_sq, log_pi)

# M-step
pi_new    = Nk / N
mu_new    = mu_acc / Nk[:, None]
sigma_sq  = sig_acc / (D * Nk)
```

### Outputs

| Tensor | Shape | Description |
|---|---|---|
| `logZ` | (N,) | Per-sample log-normaliser log Z_i |
| `Nk` | (K,) | Effective cluster counts Σ_i r_ik |
| `mu_acc` | (K, D) | Weighted sum Σ_i r_ik x_i |
| `sig_acc` | (K,) | Weighted sq. dist. Σ_i r_ik ‖x_i − μ_k‖² |

### Block sizes

The defaults `BLOCK_N=64`, `BLOCK_K=16`, `BLOCK_D=128` work for D≤128. For larger D, increase `BLOCK_D` to the next power of two ≥ D:

```python
logZ, Nk, mu_acc, sig_acc = flash_gmm_estep(X, mu, lss, lpi, BLOCK_D=256)
```

### Streaming for N > GPU memory

For datasets larger than GPU memory, feed data in chunks — the O(KD) accumulators are simply summed across chunks:

```python
Nk_total = torch.zeros(K, device='cuda')
mu_total = torch.zeros(K, D, device='cuda')
sig_total = torch.zeros(K, device='cuda')

for chunk in dataloader:  # chunks loaded from CPU/SSD
    X_chunk = chunk.cuda()
    _, Nk_c, mu_c, sig_c = flash_gmm_estep(X_chunk, mu, lss, lpi)
    Nk_total  += Nk_c
    mu_total  += mu_c
    sig_total += sig_c
    del X_chunk

# M-step on aggregated statistics
mu_new   = mu_total / Nk_total[:, None]
sigma_sq = sig_total / (D * Nk_total)
```

This was validated at **N=1B** vectors (512 GB of data) on a single A100 80GB, completing in ~28 minutes with **1,548 MB peak GPU memory**.

## Performance

Runtime of a single E-step (K=1024, D=128, A100 80GB):

| N | Flash-GMM | vs SciPy (CPU) | vs TorchGMM (GPU) |
|---|---|---|---|
| 10K | 3 ms | 766× | 32× |
| 50K | 9 ms | 1,260× | 20× |
| 100K | 18 ms | 1,458× | 23× |
| 250K | 46 ms | 1,597× | 19× |
| 500K | 84 ms | 1,571× | 20× |
| 1M | 152 ms | 1,738× | 22× |
| 10M | 1,519 ms | 1,740× | OOM |
| 50M | 35,510 ms | 1,752× | OOM |

TorchGMM runs out of memory beyond N≈1M. Flash-GMM scales to N=10⁸ on the same device.

## H100-Tuned Variants

For NVIDIA H100 on the paper benchmark workloads (K=1024, D=96–128), we ship kernels tuned for those specific shapes that run **roughly 100× faster than TorchGMM**: [`flash_gmm_h100.py`](flash_gmm_h100.py), [`flash_gmm_diag_h100.py`](flash_gmm_diag_h100.py), [`flash_gmm_full_h100.py`](flash_gmm_full_h100.py).

For shapes far from the paper benchmarks, or non-H100 GPUs, use `flash_gmm.py`.

## scikit-learn API

[`flash_gmm_sklearn.py`](flash_gmm_sklearn.py) wraps the kernels in
`FlashGaussianMixture`, a drop-in replacement for
`sklearn.mixture.GaussianMixture`. The full EM loop (E-step via the Triton
kernels, M-step in Torch) runs on the GPU; data is moved to the device once and
results come back as NumPy arrays — so integrating into an existing sklearn
workflow is a one-line change:

```python
# before
from sklearn.mixture import GaussianMixture
gm = GaussianMixture(n_components=8, covariance_type="diag").fit(X)

# after
from flash_gmm_sklearn import FlashGaussianMixture as GaussianMixture
gm = GaussianMixture(n_components=8, covariance_type="diag").fit(X)

labels = gm.predict(X)          # same attributes/methods as sklearn:
probs  = gm.predict_proba(X)    # weights_, means_, covariances_,
ll     = gm.score(X)            # precisions_cholesky_, n_iter_, bic(), aic(), sample()...
```

`covariance_type` maps to the matching kernel — `"spherical"` →
`flash_gmm_estep`, `"diag"` → `flash_gmm_estep_diag`, `"full"` →
`flash_gmm_estep_full`. (`"tied"` is not supported.) Requires `scikit-learn`
in addition to `torch`/`triton`; a CUDA GPU is required.

## Authors

- Gal Bloch (gal.bloch@ibm.com)
- Ariel Gera (ariel.gera1@ibm.com)
- Matan Orbach (matano@il.ibm.com)
- Ohad Eytan (ohad.eytan1@ibm.com)
- Assaf Toledo (assaf.toledo@ibm.com)

IBM Research

## Citation

If you use Flash-GMM in your research, please cite:

```bibtex
@article{bloch2026flashgmm,
  title     = {Flash-GMM: A Memory-Efficient Kernel for Scalable Soft Clustering},
  author    = {Bloch, Gal and Gera, Ariel and Orbach, Matan and Eytan, Ohad and Toledo, Assaf},
  journal   = {arXiv preprint arXiv:2606.10896},
  year      = {2026},
  url       = {https://arxiv.org/abs/2606.10896}
}
```

## License

Apache 2.0 — see [LICENSE](LICENSE).

Flash-GMM

Explore Similar Repositories

Was this content helpful?

Recent Feedback

Download README