장문 컨텍스트 어텐션 벤치마크: 커널 효율성에서 분산 컨텍스트 병렬 처리까지

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

초록

트랜스포머 기반 대규모 언어 모델(LLM)은 놀라운 성과를 거두었지만, 표준 어텐션 메커니즘은 시퀀스 길이에 대해 2차적인 계산 및 메모리 비용을 발생시켜 장문맥 학습의 주요 병목 현상으로 작용합니다. 선행 연구는 이 문제를 두 가지 방향으로 접근해 왔습니다: (1) 밀집 및 희소 어텐션 연산자를 가속화하는 커널 수준 최적화와, (2) 분산 어텐션 또는 컨텍스트 병렬 학습으로 불리는 모듈 수준 전략으로, 여러 장치에 걸쳐 어텐션을 확장합니다. 그러나 체계적인 평가는 여전히 제한적입니다: 연산자 수준 비교는 종종 불완전하고, 컨텍스트 병렬 전략은 일반적으로 특정 프레임워크에 종속되어 있으며 다양한 컨텍스트 간 명확한 성능 분석이 부족합니다. 이러한 격차를 해결하기 위해, 우리는 평가를 위한 모듈화되고 확장 가능한 인터페이스로 대표적인 어텐션 커널과 컨텍스트 병렬 메커니즘을 통합한 통합 벤치마크를 제안합니다. 본 벤치마크는 두 가지 핵심 차원에서 방법들을 평가합니다: (1) 효율성, 확장성 및 사용성에 큰 영향을 미치는 어텐션 마스크 패턴과, (2) 극단적인 장문맥 학습 조건에서의 성능을 결정하는 시퀀스 길이 및 분산 규모입니다. 최대 96개의 GPU로 구성된 클러스터에서의 포괄적인 실험을 통해, 우리의 벤치마크는 재현 가능한 비교를 가능하게 하고 방법별 장단점을 부각하며, 장문맥 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.