构建统一基准,评估长序列注意力的加速与分布式扩展效果。
Long-Context Attention Benchmark: From Kernel Efficiency to Distributed Context Parallelism
- 整合主流注意力核与分布式并行机制,支持可复现对比。
- 在最多96张GPU上测试,揭示不同方法在超长序列下的性能差异。
- 适合关注大模型长文本训练优化的研究者与工程师。
基于Transformer的大语言模型虽取得显著成功,但标准注意力机制随序列长度呈二次增长,导致计算与内存开销剧增,成为长序列训练的主要瓶颈。现有工作主要从两个方向应对:(1)核级优化,加速稠密与稀疏注意力算子;(2)模块级策略,即分布式注意力或上下文并行训练,实现跨多设备扩展。然而系统性评估仍不充分:算子级比较常不完整,而上下文并行策略多依赖特定框架,缺乏跨场景性能分析。为此,我们提出一个统一基准,集成代表性注意力核与上下文并行机制,提供模块化、可扩展的评估接口。该基准从两个关键维度评估方法:(1)注意力掩码模式,强烈影响效率、可扩展性与可用性;(2)序列长度与分布式规模,决定极端长序列训练下的性能表现。通过在最多96张GPU的集群上进行综合实验,该基准支持可复现比较,揭示方法间的权衡关系,并为长序列LLM训练中注意力机制的设计与部署提供实践指导。
原文摘要 · Abstract (English)
Transformer-based large language models (LLMs) have achieved remarkable success, yet their standard attention mechanism incurs quadratic computation and memory costs with respect to sequence length, posing a major bottleneck for long-context training. Prior work tackles this challenge along two directions: (1) kernel-level optimizations, which accelerate dense and sparse attention operators; and (2) module-level strategies, often referred to as distributed attention or context parallel training, which scale attention across multiple devices. However, systematic evaluation still remains limited: operator-level comparisons are often incomplete, while context parallel strategies are typically framework-specific, with unclear performance analysis across contexts. To address these gaps, we propose a unified benchmark that integrates representative attention kernels and context parallel mechanisms with a modular and extensible interface for evaluation. The benchmark evaluates methods along two critical dimensions: (1) attention mask patterns, which strongly affect efficiency, scalability, and usability, and (2) sequence length and distributed scale, which determine performance under extreme long-context training. Through comprehensive experiments on the cluster of up to 96 GPUs, our benchmark enables reproducible comparisons, highlights method-specific trade-offs, and provides practical guidance for designing and deploying attention mechanisms in long-context LLM training.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。