Speeding Up JumpReLU SAE Inference with Custom Triton Kernels (2–14× on Real SAEs)

MotivationSparse Autoencoders (SAEs) have become a central tool in mechanistic interpretability research, providing a way to decompose a model's internal activations into sparse, interpretable features. However, extracting these features often requires running the SAE over large volumes of activations across many layers and tokens. This makes SAE inference efficiency a practical bottleneck for interpretability research at scale. This post focuses on improving the inference efficiency of JumpReLU...