---
license: cc-by-nc-4.0
library_name: pytorch
tags:
- code
- triton
- pytorch
- sparse
- mixture-of-experts
- swiglu
- gpu-kernel
---

# eugen_sparse — sparse SwiGLU FFN in Triton

**Code only — no model weights.**

## Model card summary

| | |
|---|---|
| **What this repository is** | A code library (Triton GPU kernels + PyTorch wrapper) for a sparse SwiGLU feed-forward layer. **No model weights.** |
| **Developed by** | Jakub Ćwirlej |
| **Origin** | The sparse model variants of the EUGEN project (the production EUGEN 160M model has a dense FFN and does not use this kernel) |
| **Framework / hardware** | PyTorch + Triton; NVIDIA GPU with bf16 tensor cores (Ampere or newer); a slow dense reference path runs anywhere |
| **Status** | Experimental: tested in our own training runs, early observations promising; it is an open question whether it pays off when a model is trained sparse from scratch and densified at the end |
| **Intended use** | Research and experiments with sparse feed-forward layers in transformer training |
| **Out of scope** | Commercial use (licence); production training without your own validation on your GPU and model shapes |
| **Licence** | CC BY-NC 4.0 |

## In plain words

A transformer layer spends most of its compute in the feed-forward block, where every token goes through every neuron. Here the neurons are split into small groups, and for each token only the few groups that would respond most strongly are computed — the others are skipped entirely, also when computing gradients. This is a single Python file with GPU kernels written in Triton; it plugs into a PyTorch model in place of a normal SwiGLU feed-forward layer and works with `torch.compile`.

## Status

Tested experimentally; early observations are promising. It is open whether sparse training pays off over the whole cycle — training from scratch and densifying later (switching the sparse layer to a dense one). The kernel comes from the sparse model variants of the EUGEN project; the production model EUGEN 160M has a dense FFN and does not use it. This is the kernel we trained with (not a simplified demo), cleaned up for publication. Tested on one GPU type only (see Requirements). Not a validated product.

## What it does

- `make_subsample_index(d_ff, group_size, samples_per_group, seed)` — a fixed random subset of gate rows per group, used to estimate group energy cheaply. Draw once per layer and save it with the model.
- `grouped_topk_select(x, W_gate, group_size, num_active_groups, subsample_idx=None, samples_per_group=None, load_balance_bias=None)` — for every token, the indices of the `num_active_groups` groups with the largest energy `Σ|SiLU(x·W_gate)|` (optionally estimated on the subsample, optionally re-weighted by a balancing bias). Hard top-k, no gradient.
- `sparse_swiglu_ffn(x, W_gate, W_up, W_down, group_idx, group_size)` — the FFN output `Σ_selected (SiLU(x W_g) ⊙ x W_u) W_d` with gradients to `x` and all three weights. Weights in `nn.Linear` layout. One `torch.library.triton_op`: no graph breaks under `torch.compile`.
- `update_load_balance_bias(bias, freq_ema, current_freq, target_freq, gamma, bias_max, momentum)` — auxiliary-loss-free load balancing (in the style of DeepSeek-V3): call once per optimizer step.
- `reference_sparse_swiglu_ffn(...)` — a dense masked PyTorch reference with the same math (tests, CPU path).
- `kernel_supported(...)`, `plan(...)` — will the Triton path run for this shape, and how the work is split (no GPU work).

## Requirements

- NVIDIA GPU with bf16 tensor cores (Ampere or newer). Tested only on an RTX PRO 6000 (Blackwell); other cards are untested.
- PyTorch with `torch.library.triton_op` (2.6 or newer per the PyTorch docs; tested on 2.12).
- Triton 3.x (tested on 3.7.0).
- Without a GPU (or with an unsupported shape / older PyTorch) the same math runs on the slow dense reference path. On a CPU you can still check the kernels in the Triton interpreter: `TRITON_INTERPRET=1 python test_eugen_sparse.py` (needs Triton installed).

## Install

No package to install: copy `eugen_sparse.py` next to your code (it imports only `torch` and, if present, `triton`).

```bash
pip install torch triton       # versions as above
git clone https://huggingface.co/JakCwi/eugen-sparse
```

## Quick start

```python
import torch
import eugen_sparse as es

d_model, d_ff, G, K = 512, 4096, 32, 8           # 8 of 128 groups active = 6.25 % of the neurons
dev = "cuda"
W_gate = torch.randn(d_ff, d_model, device=dev, dtype=torch.bfloat16) * d_model ** -0.5
W_up   = torch.randn(d_ff, d_model, device=dev, dtype=torch.bfloat16) * d_model ** -0.5
W_down = torch.randn(d_model, d_ff, device=dev, dtype=torch.bfloat16) * d_ff ** -0.5
sub, spg = es.make_subsample_index(d_ff, G, samples_per_group=8, seed=0)   # once per layer
sub = sub.to(dev)

x = torch.randn(4096, d_model, device=dev, dtype=torch.bfloat16)
idx = es.grouped_topk_select(x, W_gate, G, K, subsample_idx=sub, samples_per_group=spg)
y = es.sparse_swiglu_ffn(x, W_gate, W_up, W_down, idx, G)                 # [4096, 512]
print(es.plan(x.shape[0], d_model, d_ff, G, K))   # kernel used? token chunks, tiles, determinism
```

A complete `nn.Module` (`EugenSparseFFN`: subsample and balancing buffers, `update_balance()` after the optimizer step, autocast, `torch.compile`) is in `example.py`; `python example.py` runs on CUDA or, slowly, on CPU.

## Parameters and how to choose them

| parameter | valid values | how to choose |
|---|---|---|
| density = `num_active_groups · group_size / d_ff` | — | 5–10 %. Our training configuration: `d_ff` 4096, `group_size` 32, 8 active groups = 6.25 %. A wide FFN at low density is preferable to a narrow one at high density. |
| `group_size` | power of two, ≥ 16 | 32 (our value). Up to 64 the fused kernels run; above 64 an unfused pair (same result, more memory traffic). Larger groups use the GPU better but make the selection coarser. |
| `num_active_groups` (K) | 1 … `d_ff / group_size` | sets the density together with `group_size`. |
| `d_ff` | multiple of `group_size`, ≥ 2 groups | need not be a power of two (1536, 3072, 3200 work). |
| `d_model` | multiple of 16 | multiples of 128 run at full tile width; smaller divisors use narrower tiles. |
| `samples_per_group` | 1 … `group_size` | 8 of 32 (a quarter of the gate) in our training. More = more accurate, more expensive selection; `= group_size` (or no subsample) = exact selection. |
| balancing: `gamma`, `bias_max`, `momentum` | — | 1e-3, 0.5, 0.99 (our values); `target_freq = num_active_groups / num_groups`; call after every `optimizer.step()`. |
| tokens per micro-batch | `tokens · K / num_groups` well above 64 | ideally hundreds of tokens per group; small micro-batches leave tiles half empty. |
| dtype | bf16 | under `torch.autocast` the function casts inputs and weights to the autocast dtype itself. fp32 works (tests) but is not the target. |
| `EUGEN_SPARSE_*` environment variables | integers, read at import | defaults = our training configuration. `EUGEN_SPARSE_DW_DET=1` forces deterministic weight gradients (slower). Change the others only with your own measurement. |

Layout that worked better for us than several dense layers: first layer dense, the following layers sparse.

## Pitfalls and limitations

- **Silent fallback.** For an unsupported shape, no Triton or an old PyTorch, the slow dense reference runs instead; on CUDA a one-time warning is printed. Check `kernel_supported(...)` / `plan(...)` once.
- **The subsample must stay fixed**: draw it once per layer, keep it in `state_dict`. A new draw per step or after resuming changes the selection.
- **The balancing bias and its EMA belong to `state_dict`** too (`persistent=True`).
- **Non-determinism**: `y` and `dX` are accumulated with fp32 atomics, so bitwise differences between runs are possible; weight gradients too when there are several token chunks (unless `EUGEN_SPARSE_DW_DET=1`). Selection and the CSR bucketing are deterministic.
- **Hard selection**: only the selected groups get weight gradients in a step; the routing itself is not learned by gradient.
- **`group_idx` must come from `grouped_topk_select`** (distinct groups per token). Hand-made indices with duplicates or out-of-range values are not checked: the kernel would count a duplicated group twice, the reference once.
- **No input validation beyond the shape rules**: `tokens · num_active_groups` must stay below 2³¹ (int32 indices).
- **Block benchmarks do not predict training speed**: measure your own training loop under `torch.compile`. Random inputs spread tokens evenly; a trained model has popular groups (the balancing bias counters that).
- **Memory**: compact activations `[tokens, K·group_size]` (gate and up) plus fp32 buffers `[tokens, d_model]` for `y` and `dX`.
- **Changing `group_size` mid-training** changes the number of groups and the buffer shapes (bias, subsample): it is a model change, not a setting.
- The example module counts group usage in its last forward call of the step; with gradient accumulation, accumulate `lb_freq_now` yourself.

## How to test

```bash
python test_eugen_sparse.py                     # GPU: everything; CPU: selection, subsample, balancing, plan
TRITON_INTERPRET=1 python test_eugen_sparse.py  # CPU: kernels vs reference in the Triton interpreter (fp32)
python speed_check.py --compile          # GPU: time dense vs sparse FFN block on YOUR card (we publish no benchmark results)
```

On a GPU the tests compare the kernel with the reference (y, dX, dW_gate, dW_up, dW_down) in fp32 and bf16, with several token chunks, run `torch.library.opcheck` and check `torch.compile` without graph breaks. How the kernels work: [TECHNICAL.md](TECHNICAL.md).

## Files

`eugen_sparse.py` (library) · `test_eugen_sparse.py` · `speed_check.py` (a timing script to run on your own GPU; no results are published) · `example.py` · `TECHNICAL.md` · `LICENSE`

## Licence

Creative Commons Attribution-NonCommercial 4.0 International (CC BY-NC 4.0), © 2026 Jakub Ćwirlej — free for non-commercial use with attribution. Full text in `LICENSE`. Dependencies (not included): PyTorch (BSD-3-Clause), Triton (MIT).

## Citation and contact

Jakub Ćwirlej, *eugen_sparse: sparse SwiGLU FFN kernel in Triton*, 2026, https://huggingface.co/JakCwi/eugen-sparse. Contact and more about the project: https://jakubcwirlej.com.