長文脈アテンションのベンチマーク：カーネル効率から分散コンテキスト並列処理へ

Long-Context Attention Benchmark: From Kernel Efficiency to Distributed Context Parallelism

October 19, 2025

著者: Tao Bu, Qiangang Wang, Bowen Zeng, Hanwen Sun, Yunpeng Huang, Chun Cao, Jingwei Xu

cs.AI

要旨

Transformerベースの大規模言語モデル（LLM）は目覚ましい成功を収めているが、標準的なアテンション機構は系列長に対して二次的な計算量とメモリコストが発生し、長文脈学習における主要なボトルネックとなっている。従来の研究はこの課題に二つの方向から取り組んできた：（1）密および疎なアテンション演算子を高速化するカーネルレベル最適化、（2）分散アテンションまたは文脈並列学習と呼ばれるモジュールレベル戦略（複数デバイス間でアテンションをスケーリングする手法）である。しかし、体系的な評価は依然として限られている：演算子レベルの比較は往々にして不完全であり、文脈並列戦略は通常フレームワーク依存で、様々な文脈における性能分析が不明確である。これらの課題を解決するため、我々は代表的なアテンションカーネルと文脈並列メカニズムを統合し、モジュール化された拡張可能な評価インターフェースを備えた統一ベンチマークを提案する。本ベンチマークは、効率性・拡張性・実用性に強く影響する（1）アテンションマスクパターンと、極端に長い文脈学習における性能を決定する（2）系列長と分散規模という二つの重要次元に沿って手法を評価する。最大96GPUクラスタでの包括的実験を通じて、本ベンチマークは再現可能な比較を実現し、手法固有のトレードオフを明らかにし、長文脈LLM学習におけるアテンション機構の設計と導入に対する実践的な指針を提供する。

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.