Measured speed in the model card and the docstring: first kernel version 1.23x faster than dense; current version 36% higher training throughput than dense (306M, d_ff 4096, 6.25% active)

README.md

@@ -23,6 +23,7 @@ tags:
 | **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 |
+| **Measured in our training** | Full model with d_model 512, d_ff 4,096 (306M parameters), 6.25% active (256 of 4,096 FFN neurons per token), against the dense model of the same FFN capacity, on one RTX PRO 6000: the **first version** of this kernel trained **1.23×** faster; the **current version** reaches **36% higher training throughput** than the dense model. A single FFN block ran 1.26× (d_ff 4,096) and 1.34× (d_ff 8,192) faster than its dense equivalent. |
 | **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 |
@@ -119,14 +120,14 @@ Layout that worked better for us than several dense layers: first layer dense, t
 ```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)
+python speed_check.py --compile          # GPU: time dense vs sparse FFN block on YOUR card (our own figures: the model card above)
 ```
 
 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`
+`eugen_sparse.py` (library) · `test_eugen_sparse.py` · `speed_check.py` (a timing script to run on your own GPU) · `example.py` · `TECHNICAL.md` · `LICENSE`
 
 ## Licence
 

eugen_sparse.py

@@ -26,8 +26,9 @@ PRACTICAL GUIDANCE
      stops paying off against a dense FFN of the same width. Our training configuration is 6.25%:
      d_ff 4096, group_size 32, 8 active groups, gate subsample of 8 rows per group.
    - Measured at 6.25%: a single FFN block ran 1.26x (d_ff 4096) and 1.34x (d_ff 8192) faster than its
-     dense equivalent; a full model (d_model 512, d_ff 4096) trained 1.23x faster than the dense model
-     of the same FFN capacity, before further kernel optimisations (+36% throughput).
+     dense equivalent. A full model (d_model 512, d_ff 4096) against the dense model of the same FFN
+     capacity: the first version of this kernel trained 1.23x faster; the current version reaches
+     36% higher training throughput than the dense model.
    - The gain grows with FFN width: prefer a wide FFN at low density over a narrow FFN at high density.
      d_ff does not have to be a power of two (e.g. 1536, 3072, 3200) — only a multiple of group_size.
    - Keep enough tokens per group in a micro-batch: num_tokens * num_active_groups / num_groups should be