APB: Accelerating Distributed Long-Context Inference by Passing Compressed Context Blocks across GPUs
Fuente:
arXiv
Saved in:
| Main Authors: | , , , , , , , , , |
|---|---|
| Format: | Preprint |
| Published: |
2025
|
| Subjects: | |
| Online Access: | |
| Tags: |
Add Tag
No Tags, Be the first to tag this record!
|
| _version_ | 1866912394063642624 |
|---|---|
| author | Huang, Yuxiang Li, Mingye Han, Xu Xiao, Chaojun Zhao, Weilin Ao, Sun Zhou, Hao Zhou, Jie Liu, Zhiyuan Sun, Maosong |
| author_facet | Huang, Yuxiang Li, Mingye Han, Xu Xiao, Chaojun Zhao, Weilin Ao, Sun Zhou, Hao Zhou, Jie Liu, Zhiyuan Sun, Maosong |
| contents | While long-context inference is crucial for advancing large language model (LLM) applications, its prefill speed remains a significant bottleneck. Current approaches, including sequence parallelism strategies and compute reduction through approximate attention mechanisms, still fall short of delivering optimal inference efficiency. This hinders scaling the inputs to longer sequences and processing long-context queries in a timely manner. To address this, we introduce APB, an efficient long-context inference framework that leverages multi-host approximate attention to enhance prefill speed by reducing compute and enhancing parallelism simultaneously. APB introduces a communication mechanism for essential key-value pairs within a sequence parallelism framework, enabling a faster inference speed while maintaining task performance. We implement APB by incorporating a tailored FlashAttn kernel alongside optimized distribution strategies, supporting diverse models and parallelism configurations. APB achieves speedups of up to 9.2x, 4.2x, and 1.6x compared with FlashAttn, RingAttn, and StarAttn, respectively, without any observable task performance degradation. We provide the implementation and experiment code of APB in https://github.com/thunlp/APB. |
| format | Preprint |
| id |
arxiv_https___arxiv_org_abs_2502_12085 |
| institution | arXiv |
| publishDate | 2025 |
| record_format | arxiv |
| spellingShingle | APB: Accelerating Distributed Long-Context Inference by Passing Compressed Context Blocks across GPUs Huang, Yuxiang Li, Mingye Han, Xu Xiao, Chaojun Zhao, Weilin Ao, Sun Zhou, Hao Zhou, Jie Liu, Zhiyuan Sun, Maosong Machine Learning Computation and Language While long-context inference is crucial for advancing large language model (LLM) applications, its prefill speed remains a significant bottleneck. Current approaches, including sequence parallelism strategies and compute reduction through approximate attention mechanisms, still fall short of delivering optimal inference efficiency. This hinders scaling the inputs to longer sequences and processing long-context queries in a timely manner. To address this, we introduce APB, an efficient long-context inference framework that leverages multi-host approximate attention to enhance prefill speed by reducing compute and enhancing parallelism simultaneously. APB introduces a communication mechanism for essential key-value pairs within a sequence parallelism framework, enabling a faster inference speed while maintaining task performance. We implement APB by incorporating a tailored FlashAttn kernel alongside optimized distribution strategies, supporting diverse models and parallelism configurations. APB achieves speedups of up to 9.2x, 4.2x, and 1.6x compared with FlashAttn, RingAttn, and StarAttn, respectively, without any observable task performance degradation. We provide the implementation and experiment code of APB in https://github.com/thunlp/APB. |
| title | APB: Accelerating Distributed Long-Context Inference by Passing Compressed Context Blocks across GPUs |
| topic | Machine Learning Computation and Language |
| url | https://arxiv.org/abs/2502.12085 |