Long-Context Attention Benchmark: From Kernel Efficiency to Distributed Context Parallelism
This work addresses the bottleneck of long-context training for LLM developers and researchers, offering a systematic benchmark to guide method selection, but it is incremental as it builds on existing optimizations without introducing new methods.
The authors tackled the problem of evaluating attention mechanisms for long-context training in large language models, which suffer from quadratic costs, by proposing a unified benchmark that integrates kernel-level optimizations and context parallel strategies. The result is a comprehensive evaluation on up to 96 GPUs, providing reproducible comparisons and practical guidance for method-specific trade-offs.
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.