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

TL;DR

The paper tackles the inefficiency of attention mechanisms for ultra-long contexts by introducing LongCA-bench, a unified benchmarking framework that jointly assesses kernel-level attention operators and context-parallel distributed mechanisms. It provides modular components for dense and sparse kernels, varied mask patterns, and scalable distributed strategies, enabling reproducible comparisons up to 512K tokens on 96 GPUs. Through large-scale experiments, the authors reveal complementary trade-offs: hardware-optimized dense kernels achieve high throughput but memory bottlenecks persist, while block-sparse kernels offer efficiency with limitations in backward support and mask flexibility. The benchmark also exposes design considerations for context parallelism, showing that hybrid architectures can improve scalability and stability, thereby guiding practical deployment of long-context training systems in real-world LLM workloads.

Abstract

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.

Figures (39)

• Figure 1: The architecture of LongCA benchmark
• Figure 2: Attention mask patterns
• Figure 3: Forward TFLOPs of dense kernels with different masks (8K length)
• Figure 4: Performance results (TFLOPs) of sparse kernels with a 50% sparsity ratio
• Figure 5: Forward TFLOPs of Context Parallel Attention on FULL DOCUMENT.