File size: 1,193 Bytes
7cf627f | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 | # Sparse Transformer: Experiment Suite + Triton Kernels
Comprehensive experiment infrastructure for the Chunked Sparse Backward Pass paper.
## Files
| File | Description |
|------|-------------|
| `triton_sparse.py` | Triton-fused sparse backward kernels (dW, dX, dBias) + Python-loop baseline + correctness tests + microbenchmark |
| `e2e_full.py` | End-to-end training benchmark: Dense vs PyLoop vs Triton at d_model ∈ {512, 1024, 2048} |
| `full_experiments.py` | 7-experiment ablation suite (baselines, predictor accuracy, chunk ablation, compute-matched, exploration, attention sparsification, sparsity sweep) |
| `analyze_results.py` | Publication figure generator (matplotlib) |
## Quick Start
```bash
pip install torch triton tiktoken matplotlib numpy
# Correctness test + microbenchmark
python triton_sparse.py
# End-to-end training (needs ≥24GB GPU for d=2048)
python e2e_full.py
# Full ablation suite (7 experiments, ~4-6 hours on A10G)
python full_experiments.py --experiment all --device cuda --steps 2000 --seeds "42,123,456"
# Single experiment
python full_experiments.py --experiment baselines --device cuda
```
## Results
See `RESULTS.md` for collected tables.
|