Sparse Attention Unified

SSA sparse attention training with SOCKET and Soft-SAAP backends

Source Code

• sparse-attn-unified

Sparse Attention Unified is an SSA (Sparse Sparse Attention) training framework with pluggable SOCKET (soft-LSH) and Soft-SAAP (Gumbel-Softmax clusters) sparse backends. It targets efficient training of sparse attention models such as Llama 3.2 1B on a single H100 GPU.

During training, both full attention (FA) and sparse attention (SA) streams run at every layer; an alignment loss keeps sparse outputs close to full attention. At inference, only the sparse backend runs, enabling faster decoding while preserving model quality.

The repository includes HPC setup scripts for the SJSU CoE cluster, smoke tests, SLURM job templates, and configuration files for switching between SOCKET and SAAP backends.

Contributors

To be added.

Publications

To be added.

← Back to all projects