This repository hosts the official source code of SharePrefill, a novel sparse attention method that highly preserves accuracy while accelerating the prefilling phase during long-context inference of LLMs.
- Python 3.10
- PyTorch == 2.5.1
- Transformers == 4.51.3
- Triton ≥ 3.1.0
- Optional:
flash-attnfor FlashAttention kernels - Tested on NVIDIA A100 GPUs with 80GB VRAM
cd SharePrefill
pip install -e .model = AutoModelForCausalLM.from_pretrained(
model_name,
torch_dtype="auto",
device_map="auto",
)
shareprefill = SharePrefill(
"shareprefill",
model_name=model_name,
block_size=128,
min_budget=1024,
similarity_threshold=0.5,
pattern_threshold=0.95,
)
model = shareprefill(model)Evaluate long-context accuracy with RULER.
cd SharePrefill/eval/ruler
pip install -r requirements.txt
python3 ./download_nltk.py
bash run_ruler.shcd SharePrefill/eval/infinite_bench
pip install -r requirements.txt
bash run_infinitebench.shcd SharePrefill/eval/efficiency
wget https://raw.githubusercontent.com/FranxYao/chain-of-thought-hub/main/gsm8k/lib_prompt/prompt_hardest.txt
bash run_latency.shIf you find SharePrefill helpful in your research, please cite:
@article{peng2025accelerating,
title={Accelerating Prefilling for Long-Context LLMs via Sparse Pattern Sharing},
author={Peng, Dan and Fu, Zhihui and Ye, Zewen and Song, Zhuoran and Wang, Jun},
journal={arXiv preprint arXiv:2505.19578},
year={2025}
}
SharePrefill is released under the Apache License 2.0. Refer to the repository’s LICENSE file for the full terms. This project incorporates code snippets derived from MInference and FlexPrefill, and reuses evaluation assets from RULER. We thank their authors and maintainers for sharing their work. Third-party notices are listed in NOTICE and full license texts are under licenses/.