JakCwi/eugen-sparse

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).

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

Quick start

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

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.

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.