PaperlayerView exact revision on arXiv

Source-linked AI summary

Efficient Long-context Language Model Training by Core Attention Disaggregation

Yonghao Zhuang, Junda Chen, Bo Pang, Yi Gu, Yibo Zhu, Yimin Jiang, Ion Stoica, Eric Xing, Hao Zhang

arXiv:2510.18121v1cs.LGcs.DC

TL;DR

Packed long-document batches can create attention-compute imbalance because attention scales quadratically with sequence length. CAD separates core attention into token-level tasks scheduled on dedicated attention servers. DistCA reports up to 1.35x throughput speedup over state-of-the-art training systems.

  • Problem

    Packed chunks with equal token counts can require different attention FLOPs, creating load imbalance because attention compute grows quadratically with sequence length.

    Source: §1 · p. 1

  • Method

    CAD partitions core attention into token-level document-shard tasks and dynamically batches them for scheduling on a dedicated pool of attention servers.

    Source: §1 · p. 2

  • Results

    Up to 1.35x throughput speedup over state-of-the-art training systems is reported for DistCA.

    Source: §9 · pp. 10–11

  • Takeaways & Limitations

    CAD eliminates data- and pipeline-parallel stragglers while achieving near-perfect load balancing across attention servers.

    Source: §1 · p. 2

    The scheduler restricts each core-attention task to a Q shard paired with the full K,V context, limiting shard-placement flexibility.

    Source: §8 · p. 10

Abstract

from arXiv · show

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.

1 INTRODUCTION

Long-context training faces attention-driven load imbalance because attention grows quadratically while other model computation grows approximately linearly. CAD separates and dynamically schedules core-attention tasks; DistCA reports higher end-to-end throughput.

1 INTRODUCTION

  • Applications require 100K–1M-token contexts, but packed documents create variable quadratic attention work, causing DP and PP stragglers and reported 1.34–1.44x slowdowns at modest context lengths.Source: §1 · pp. 1–2
  • Matching attention FLOPs by assigning more tokens to devices processing shorter documents raises activation memory 3× in the 4×1K example, while context parallelism adds KV communication and cannot eliminate PP stragglers.Source: §1 · p. 2
  • CAD separates parameter-free core attention into arbitrary token shards, dynamically rebatching them across documents into fused kernels to balance compute independently of other model layers.Core attention stores only small per-row softmax statistics, and modern kernels retain throughput when fused shards come from different documents.Source: §1 · p. 2
  • DistCA time-shares GPUs between attention servers and model computation, overlaps communication with computation through ping-pong batches, and schedules document sharding to balance workload and communication.The paper lists in-place GPU time sharing, ping-pong overlap, and a workload-balanced scheduler among DistCA’s key optimizations.Source: §1 · pp. 2–3
  • DistCA improves end-to-end throughput by up to 1.35x versus existing systems at up to 512 H200 GPUs and 512k-token contexts, with near-linear weak scaling.The study also reports that CAD communication can be fully hidden and attention balanced near-perfectly.Source: §1 · p. 2

2 BACKGROUND

The background distinguishes token-wise context-independent layers from stateless core attention, then explains how uneven workloads can disrupt training parallelism.

2.1 LLM Architecture

  • Context-independent layers operate token-wise, so their compute and activation memory scale approximately linearly with token count.They include QKV and output projections, FFN, layernorm, and position embedding; GEMM operators dominate their cost.Source: §2.1 · p. 3
  • Core attention (CA) is the only context-dependent layer, computing softmax(QK^T)V without trainable parameters or significant saved intermediates.Training kernels avoid storing the quadratic attention matrix and recompute it during backward.Source: §2.1 · p. 3

2.2 LLM Training Parallelization

  • Uneven packed-data attention workloads create data-parallel stragglers, while variable-length microbatches imbalance pipeline stages and propagate bubbles downstream.Global data resharding cannot fix pipeline imbalance because each stage processes different model layers.Source: §2.2 · p. 3
  • Tensor parallelism shards layers and balances compute and memory, but per-layer communication makes scaling beyond one node unaffordable.Attention is sharded across heads, and all tensor-parallel devices process the same batch.Source: §2.2 · p. 3
  • Causal attention makes later sequence shards more compute-heavy, so naive context-parallel slicing creates load imbalance.Head-tail shard assignment mitigates this imbalance; with document packing, context-parallel shards are assigned per document (§3.2).Source: §2.2 · p. 3

3 CHALLENGE AND MOTIVATION

Section 3 explains that quadratic attention compute and linear activation memory make packed workloads difficult to balance simultaneously; existing remedies introduce trade-offs that motivate separating core attention from context-independent layers.

3.1 Load Imbalance in Compute and Memory

  • Compute scales as FLOPs(l) = αl2 + βl while activation memory scales as M(l) = γl, so equal token counts do not ensure equal attention work or simultaneous compute-memory balance.The quadratic term represents core attention FLOPs; the linear terms aggregate context-independent layer FLOPs and activation memory.Source: §3.1 · pp. 3–4

3.2 Problems of Existing Methods

  • At 512K maximum length, compute-balanced variable-length chunking needs 1.08 - 1.17x more activation memory on some ranks; average idle time reaches 19% at DP=4 and 55% at DP=8.Under memory pressure, moving sequences cannot fully equalize attention compute, making variable-length chunking ineffective.Source: §3.2 · p. 4, Figure 4 · p. 4
  • Per-document CP balances per-document compute and memory, but tiny shards underfill tiles; all-gather latency rises from 3% at 2 nodes to nearly 40% at 32, while KV memory grows from 3% at 2 nodes to almost 30% at 16.The increasing communication and memory overhead, alongside reduced kernel throughput for documents shorter than 128 tokens, limits scalability.Source: §3.2 · p. 4, Figure 3 · p. 4, Figure 5 · p. 5
  • On a 64-GPU 512K-token workload, combining variable-length chunking with per-document CP trades improved balance at higher CP degree for lower throughput and OOM risk; higher DP causes severe imbalance.The trade-off becomes more acute at larger scale and longer context.Source: §3.2 · p. 4, Figure 6 · p. 5

3.3 Motivation for Core Attention Disaggregation

  • CAD separates quadratic core-attention compute from linear context-independent work, enabling independent attention-compute scheduling and compute-memory balancing for the remaining layers.Because core attention is stateless, its scheduling can focus on equalizing compute across a shared pool of attention servers.Source: §3.3 · p. 4
  • Core attention is divisible into token-level shards, and fused variable-length kernels retain high throughput when aggregate work fills a tile; FA2 profiling supports shards of at least 128 tokens.Shards shorter than 128 tokens are padded, wasting the compute of their assigned thread blocks.Source: §3.3 · p. 5, Figure 5 · p. 5
  • CAD limits communication through causal-mask-aware KV transfers, distributing communication-heavy later shards, arbitrary sharding, and ping-pong overlap of one batch's transfer with another's computation.For Llama-3-34B with InfiniBand bandwidth, documents can be partitioned into up to 31 shards without communication overhead; the upper bound increases for larger models.Source: §3.3 · p. 5, Figure 7 · p. 5

4 METHOD

DistCA implements core attention disaggregation with a runtime that alternates context-independent computation and attention-server execution, plus a scheduler that shards documents and assigns work across devices. Its design combines dynamically batched attention tasks with communication-aware load balancing.

4.1 Runtime

  • DistCA splits document attention into query-shard tasks, assigns them to attention servers, batches each server’s tasks into one kernel, and returns results to the originating GPUs.The runtime alternates core attention with context-independent layers; the scheduler determines sharding and task placement.Source: §4 · p. 5, §4.1 · p. 6
  • In-place attention servers let GPUs switch between core attention and context-independent work, while ping-pong nano-batches overlap communication with computation.The schedule also overlaps intra-node tensor-parallel communication with inter-node communication from attention disaggregation.Source: §4.1 · p. 6
  • CAD supports pipeline parallelism by balancing weightless attention tasks across servers and keeping context-independent work balanced across stages.It defers selected backward microbatches into pipeline bubbles without adding ticks and uses warmup and drain-down idle GPU time for attention tasks; this supports 1F1B and interleaved 1F1B schedules.Source: §4.1 · p. 6

4.2 Communication-Aware Greedy Scheduling

  • The communication-aware scheduler profiles shard costs and seeks to balance FLOPs across attention servers while minimizing Q, K, and V transfer volume.Scheduling items are complete documents or block-aligned shards; the scheduler decides whether to split them and assigns each to a server, per microbatch or logical tick.Source: §4.2 · pp. 6–7
  • To reduce load gaps, the scheduler migrates or splits items according to a cost-benefit score that accounts for transferable FLOPs and communication cost.It first determines ideal per-server load, then selects candidate shards to move between surplus and deficit servers.Source: §4.2 · p. 7
  • Scheduling stops when server loads fall within a tolerance of the target or further moves fail to improve the migration score enough to justify communication.This termination rule limits unnecessary transfers from insignificant migrations.Source: §4.2 · p. 7

5 IMPLEMENTATIONS

DistCA adds attention-specific dispatch and execution components while reusing Megatron-LM’s existing model and training infrastructure.

5 IMPLEMENTATIONS

  • DistCA was implemented with 2K lines of Python, plus a 1K-line CUDA/C++ All-to-All kernel for efficient attention-server input and output dispatch, following Lei et al. (2025).The communication uses NVSHMEM (NVIDIA Corporation, 2025).Source: §5 · p. 7
  • DistCA integrates with Megatron-LM to reuse its token-independent layers, model architecture, 4D parallelization, and end-to-end training pipeline.The integration adds 1K lines of Python for the custom attention layer, ping-pong computation, and pipeline parallelism.Source: §5 · p. 7

6 EXPERIMENTS

DistCA generally outperforms WLB-ideal in 3D and 4D parallel experiments, while its communication overlap and scheduler tolerance shape overhead and memory–latency tradeoffs. Memory fragmentation limits performance in 34B 4D experiments.

6.1 Setup

  • Evaluation uses LLaMA 8B and 34B on DGX H200 nodes, averaging throughput over 30 batches from synthetic Pretrain and ProLong distributions across varying MaxDocLen.The total token count is set to the baseline’s memory limit because the baseline runs out of memory before DistCA.Source: §6.1 · pp. 7–8

6.2 End-to-End Experiment

  • Without PP, DistCA outperforms WLB-ideal by 1.07-1.20x on Pretrain and 1.05-1.12x on ProLong.DistCA dispatches only a subset of tokens and uses ping-pong execution to hide communication, with flexible token-level scheduling for load balance.Source: §6.2 · p. 8, Figure 9 · p. 8
  • Across 8B configurations, DistCA achieves 1.15x-1.30x speedup on Pretrain and 1.10x-1.35x on ProLong; for 34B, it reaches up to 1.15x on Pretrain and 1.25x on ProLong.In 4D parallelism with PP, DistCA generally outperforms WLB-ideal and scales more favorably across configurations.Source: Figure 10 · p. 9, §6.2 · p. 9
  • In 34B 4D experiments, variable-shaped requests trigger allocator fragmentation and frequent PyTorch garbage collection, adding CPU launch delays and degrading performance.The paper identifies static memory allocation and CUDA Graphs as future approaches to address this overhead.Source: §6.2 · p. 9

6.3 Ablation study

  • DistCA’s latency nearly matches the 1-byte “Signal” baseline, avoiding the 10-17% higher latency of “Single Stream” except for 8B on 8 nodes.The exception occurs when the compute workload is too small to fully hide communication.Source: §6.3 · p. 9, Figure 11 · p. 9
  • For 8B, latency stays largely unchanged at tolerance 0-0.20; for 34B, tolerance below 0.10 raises latency, while overly large tolerance raises it through load imbalance.For 34B, a tolerance below 0.10 increases communication volume so it can no longer be hidden; large tolerance increases latency roughly linearly.Source: §6.3 · pp. 9–10
  • Tuning tolerance from 0 to 0.15 cuts memory requirement by 20-25% in most cases while leaving average duration nearly unchanged or sometimes better; beyond that, iteration latency rises.Communication size remains relatively stable beyond the 0.15 tolerance point.Source: §6.3 · p. 10, Figure 12 · p. 10

7 RELATED WORKS

Prior work addresses long-context load imbalance through context-parallel and workload-balancing strategies, while model disaggregation has mainly targeted inference. CAD instead uses token-granularity attention sharding for training.

Load imbalance in long-context training

  • CAD uses token-granularity variable-length sharding to balance load more precisely and reduce communication compared with context-parallel strategies.It avoids sharding short documents and shards long documents with communication efficiency in mind.Source: §7 · p. 10
  • Prior context-parallel approaches include dynamic CP degree selection in FlexSP (Wang et al., 2025b), per-document dynamic CP degree, and overlapping communication with attention (Liu et al., 2024a; Chen et al., 2025).FlexSP uses an NP-hard ILP solver, limiting scalability, and targets the narrower FSDP setting.Source: §7 · p. 10
  • WLB-LLM (Wang et al., 2025c) uses variable-size data chunks and an imbalanced MLP to compensate for attention workload imbalance.Because activation memory is tied to MLP computation, equalizing total FLOPs can create memory imbalance; as context grows, the MLP share shrinks, reducing its ability to offset attention imbalance.Source: §7 · p. 10

Model disaggregation

  • Attention and FFN disaggregation has been used in inference, including MegaScale-Infer (Zhu et al., 2025) and concurrent work (Wang et al., 2025a), primarily for MoE models.Their per-layer transfers can merge with MoE token-routing communication; dense models would incur extra overhead, whereas CAD targets throughput-oriented training and reports benefits independent of model architecture.Source: §7 · p. 10

8 LIMITATIONS

DistCA’s current scheduler constrains each core-attention task to a Q shard paired with the full K, V context, and its communication model can overestimate transfer volume. The paper identifies opportunities to reduce compute time and improve fault tolerance and performance isolation, leaving the latter for future work.

8 LIMITATIONS

  • If memory demand is met, dedicating more GPUs to attention could reduce compute time while preserving load balance and low communication overhead.The authors also propose that a dedicated pool could improve fault tolerance and performance isolation, which they leave for future work.Source: §8 · p. 10
  • Using Q shards with only sub-ranges of the K, V context could give the scheduler more flexibility.Currently, each core-attention task uses a Q shard with the full K, V context.Source: §8 · p. 10
  • The communication model can overestimate bytes and produce non-minimal transfers by assuming all tokens move and ignoring destination-resident K, V.Source: §8 · p. 10

9 CONCLUSION

Core attention disaggregation separates stateless, token-composable core attention from the rest of model training, enabling independent scaling and scheduling. DistCA implements this approach and improves throughput by up to 1.35x over state-of-the-art training systems, with increasing advantages at larger scales.

9 CONCLUSION

  • Core attention disaggregation separates core attention from the rest of the model, using its statelessness and token-granularity composability for independent scaling and scheduling with minimal communication overhead.Source: §9 · p. 10
  • DistCA uses a workload-aware scheduler to balance computation while minimizing communication, and a ping-pong execution scheme to hide dispatch latency.Source: §9 · p. 10
  • Up to 1.35x speedup over state-of-the-art training systems, with increasing advantages at larger scales, is achieved by DistCA.End-to-end evaluation reports the speedup.Source: §9 · pp. 10–11

A UPPER BOUND FOR CORE ATTENTION SERVER MAX PARTITION SIZE

The section derives a communication-overlap upper bound on core attention server shards and concludes that the bound increases for larger models because context-independent layer compute time grows quadratically with hidden size.

  • When communication is fully overlapped with computation, the shard count is bounded by s ≤ 2(tB − hq)/hkv − 1.The bound follows from t · l ≥ l · (hq + hkv(s + 1)/2)/B, where B is network bandwidth and t is the time to compute a token’s Context-Independent layer.Source: A UPPER BOUND FOR CORE ATTENTION… · p. 13
  • The paper uses Llama-34B to analyze theoretical maximum partition count, assuming 50 GB/s InfiniBand bandwidth and Context Independent layer MFU of 50% of an H200 node (990 TFLOPs for FP16).Source: A UPPER BOUND FOR CORE ATTENTION… · p. 13
  • The token’s Context Independent layer requires 2h(2h + hkv + 3i) FLOPs, covering attention projections and a gated two-layer MLP.The factor of 2 counts a multiplication and an addition per matrix-multiplication element; the MLP uses three mappings between hidden and intermediate dimensions.Source: A UPPER BOUND FOR CORE ATTENTION… · p. 13
  • Because context-independent layer compute time scales quadratically with hidden size hq, the upper bound on s increases for larger models.Source: A UPPER BOUND FOR CORE ATTENTION… · p. 13

B COMMUNICATION OVERHEAD FUNCTION

The section models communication cost as a function of shard size and derives configurations that minimize it, then adjusts the model for head-tail context parallelism.

B COMMUNICATION OVERHEAD FUNCTION

  • The analysis selects shard sizes and configurations to minimize communication under the stated constraints.For one case, the optimum corresponds to the upper bound on n_max.Source: B COMMUNICATION OVERHEAD FUNCTION · pp. 13–14
  • Because shard MFU diverges, FLOP-based timing estimates are imprecise, but profiling finds them accurate for head-tail context parallelism.The authors retain head-tail context parallelism and leave more precise modeling for future work.Source: B COMMUNICATION OVERHEAD FUNCTION · p. 13
  • In head-tail sharding, the first token range needs KV states through j, while the last needs states at L_doc−i, with i+n_q=j=n_kv.These dependencies determine the modified communication accounting.Source: B COMMUNICATION OVERHEAD FUNCTION · p. 14

Related papers

Loading 2510.18121v1…