Sparser, Faster, Lighter Transformer Language Models

This repository provides custom CUDA kernels and sparse training code to significantly improve the throughput and reduce memory requirements of Transformer language models during inference and training. Targeting researchers and engineers working with large models, particularly on NVIDIA H100 GPUs, it enables sparser, faster, and lighter LLMs through optimized kernels and the TwELL packing format.

How It Works

The core innovation lies in custom CUDA kernels, specifically designed for NVIDIA H100 GPUs, that leverage the TwELL packing format to efficiently process sparse model weights. This approach exploits model sparsity to accelerate computations and decrease memory footprints, offering a performance boost over standard dense implementations. The project also includes code for sparse training and hybrid kernels that combine TwELL with other techniques for further optimization.

Quick Start & Requirements

Installation requires cloning the repository and running bash scripts/install.sh. A CUDA 12.8+ environment is mandatory. Pretrained sparse checkpoints are available on the Hugging Face Hub (e.g., SakanaAI/SparseLM1.5B). Links to the associated paper and blog are mentioned but not directly provided.

Highlighted Details

• Offers pretrained sparse checkpoints for models ranging from 0.5B to 2B parameters.
• Includes benchmarking scripts (benchmark_inference.py) to compare custom kernels against Hugging Face PyTorch implementations.
• Features a twell-flex kernel variant for potentially improved performance on non-uniform sparsity patterns.
• Supports optional GPU energy measurement during inference benchmarks.

Maintenance & Community

The README does not detail specific contributors, community channels (like Discord or Slack), or a public roadmap beyond listing completed features.

Licensing & Compatibility

No license information is provided within the README. This omission requires further investigation for commercial use or integration into closed-source projects.

Limitations & Caveats

A strict requirement for CUDA 12.8+ may limit compatibility. The kernels are optimized for H100 GPUs, suggesting performance may be suboptimal on other hardware. The absence of explicit licensing information is a significant caveat for adoption.