Efficient Long-context Language Model Training by Core Attention Disaggregation

2025/10/20 by Yonghao Zhuang, Zhuang, Yonghao, Chen, Junda +14 · 1 citation
Computer Science · #Big Data and Digital Economy #Advanced Neural Network Applications #Natural Language Processing Techniques

paper · pdf · doi:10.48550/arxiv.2510.18121

Abstract

We present core attention disaggregation (CAD), a technique that improves long-context large language model training by decoupling the core attention computation, softmax(QK^T)V, from the rest of the model and executing it on a separate pool of devices. In existing systems, core attention is colocated with other layers; at long context lengths, its quadratic compute growth compared to the near-linear growth of other components causes load imbalance and stragglers across data and pipeline parallel groups. CAD is enabled by two observations. First, core attention is stateless: it has no trainable parameters and only minimal transient data, so balancing reduces to scheduling compute-bound tasks. Second, it is composable: modern attention kernels retain high efficiency when processing fused batches of token-level shards with arbitrary lengths. CAD partitions core attention into token-level tasks and dispatches them to dedicated attention servers, which dynamically rebatch tasks to equalize compute without sacrificing kernel efficiency. We implement CAD in a system called DistCA, which uses a ping-pong execution scheme to fully overlap communication with computation and in-place execution on attention servers to reduce memory use. On 512 H200 GPUs and context lengths up to 512k tokens, DistCA improves end-to-end training throughput by up to 1.35x, eliminates data and pipeline parallel stragglers, and achieves near-perfect compute and memory balance.

Citations

• Zeppelin: Balancing Variable-length Workloads in Data Parallel Large Model Training
• Understanding Stragglers in Large Model Training Using What-if Analysis
• WLB-LLM: Workload-Balanced 4D Parallelism for Large Language Model Training
• FlexSP: Accelerating Large Language Model Training via Flexible Sequence Parallelism
• Packing Analysis: Packing Is More Appropriate for Large Models or Datasets in Supervised Fine-tuning
• How to Train Long-Context Language Models (Effectively)
• Data Engineering for Scaling Language Models to 128K Context
• Splitwise: Efficient generative LLM inference using phase splitting
• Ring Attention with Blockwise Transformers for Near-Infinite Context
• DeepSpeed Ulysses: System Optimizations for Enabling Training of Extreme Long Sequence Transformer Models
• RepoBench: Benchmarking Repository-Level Code Auto-Completion Systems
• FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness
• Scaling Language Models: Methods, Analysis & Insights from Training Gopher
• Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism
• FAST: An Efficient Scheduler for All-to-All GPU Communication

Cited by

• Libra: Taming Attention Workload Skew in Long-Context LLM Training with Bounded Sequence Pool

Related

• Long-Context Attention Benchmark: From Kernel Efficiency to Distributed Context Parallelism
• Native Sparse Attention: Hardware-Aligned and Natively Trainable Sparse Attention
• MTraining: Distributed Dynamic Sparse Attention for Efficient Ultra-Long Context Training
• A Training-free Sub-quadratic Cost Transformer Model Serving Framework With Hierarchically Pruned Attention
• A Little Goes a Long Way: Efficient Long Context Training and Inference with Partial Contexts