Mixture of Attention Spans: Optimizing LLM Inference Efficiency with Heterogeneous Sliding-Window Lengths

Fuente: arXiv
Gespeichert in:
Bibliographische Detailangaben
Hauptverfasser: Fu, Tianyu, Huang, Haofeng, Ning, Xuefei, Zhang, Genghan, Chen, Boju, Wu, Tianqi, Wang, Hongyi, Huang, Zixiao, Li, Shiyao, Yan, Shengen, Dai, Guohao, Yang, Huazhong, Wang, Yu
Format: Preprint
Veröffentlicht: 2024
Schlagworte:
Online-Zugang:
Tags: Tag hinzufügen
Keine Tags, Fügen Sie den ersten Tag hinzu!
_version_ 1866915635461619712
author Fu, Tianyu
Huang, Haofeng
Ning, Xuefei
Zhang, Genghan
Chen, Boju
Wu, Tianqi
Wang, Hongyi
Huang, Zixiao
Li, Shiyao
Yan, Shengen
Dai, Guohao
Yang, Huazhong
Wang, Yu
author_facet Fu, Tianyu
Huang, Haofeng
Ning, Xuefei
Zhang, Genghan
Chen, Boju
Wu, Tianqi
Wang, Hongyi
Huang, Zixiao
Li, Shiyao
Yan, Shengen
Dai, Guohao
Yang, Huazhong
Wang, Yu
contents Sliding-window attention offers a hardware-efficient solution to the memory and throughput challenges of Large Language Models (LLMs) in long-context scenarios. Existing methods typically employ a single window length across all attention heads and input sizes. However, this uniform approach fails to capture the heterogeneous attention patterns inherent in LLMs, ignoring their distinct accuracy-latency trade-offs. To address this challenge, we propose *Mixture of Attention Spans* (MoA), which automatically tailors distinct sliding-window length configurations to different heads and layers. MoA constructs and navigates a search space of various window lengths and their scaling rules relative to input sizes. It profiles the model, evaluates potential configurations, and pinpoints the optimal length configurations for each head. MoA adapts to varying input sizes, revealing that some attention heads expand their focus to accommodate longer inputs, while other heads consistently concentrate on fixed-length local contexts. Experiments show that MoA increases the effective context length by 3.9x with the same average sliding-window length, boosting retrieval accuracy by 1.5-7.1x over the uniform-window baseline across Vicuna-{7B, 13B} and Llama3-{8B, 70B} models. Moreover, MoA narrows the performance gap with full attention, reducing the maximum relative performance drop from 9%-36% to within 5% across three long-context understanding benchmarks. MoA achieves a 1.2-1.4x GPU memory reduction, boosting decode throughput by 6.6-8.2x and 1.7-1.9x over FlashAttention2 and vLLM, with minimal performance impact. Our code is available at: https://github.com/thu-nics/MoA
format Preprint
id arxiv_https___arxiv_org_abs_2406_14909
institution arXiv
publishDate 2024
record_format arxiv
spellingShingle Mixture of Attention Spans: Optimizing LLM Inference Efficiency with Heterogeneous Sliding-Window Lengths
Fu, Tianyu
Huang, Haofeng
Ning, Xuefei
Zhang, Genghan
Chen, Boju
Wu, Tianqi
Wang, Hongyi
Huang, Zixiao
Li, Shiyao
Yan, Shengen
Dai, Guohao
Yang, Huazhong
Wang, Yu
Machine Learning
Artificial Intelligence
Computation and Language
I.2.7
Sliding-window attention offers a hardware-efficient solution to the memory and throughput challenges of Large Language Models (LLMs) in long-context scenarios. Existing methods typically employ a single window length across all attention heads and input sizes. However, this uniform approach fails to capture the heterogeneous attention patterns inherent in LLMs, ignoring their distinct accuracy-latency trade-offs. To address this challenge, we propose *Mixture of Attention Spans* (MoA), which automatically tailors distinct sliding-window length configurations to different heads and layers. MoA constructs and navigates a search space of various window lengths and their scaling rules relative to input sizes. It profiles the model, evaluates potential configurations, and pinpoints the optimal length configurations for each head. MoA adapts to varying input sizes, revealing that some attention heads expand their focus to accommodate longer inputs, while other heads consistently concentrate on fixed-length local contexts. Experiments show that MoA increases the effective context length by 3.9x with the same average sliding-window length, boosting retrieval accuracy by 1.5-7.1x over the uniform-window baseline across Vicuna-{7B, 13B} and Llama3-{8B, 70B} models. Moreover, MoA narrows the performance gap with full attention, reducing the maximum relative performance drop from 9%-36% to within 5% across three long-context understanding benchmarks. MoA achieves a 1.2-1.4x GPU memory reduction, boosting decode throughput by 6.6-8.2x and 1.7-1.9x over FlashAttention2 and vLLM, with minimal performance impact. Our code is available at: https://github.com/thu-nics/MoA
title Mixture of Attention Spans: Optimizing LLM Inference Efficiency with Heterogeneous Sliding-Window Lengths
topic Machine Learning
Artificial Intelligence
Computation and Language
I.2.7
url https://arxiv.org/abs/2406.14909