Efficient Long-context Language Model Training by Core Attention
  Disaggregation
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.
68f90580b9b2e4ae04673658
Yonghao Zhuang
claimed_verified
2025-10-24T16:13:20.891Z
64ab987073790912c7ad79e6
/avatars/bfa00418c46da0fd1e522d6561944ec9.svg
Yonghao Zhuang
user
ZYHowell
68f90580b9b2e4ae04673659
Junda Chen
claimed_verified
2025-10-24T16:13:24.678Z
643839d9581e6bf0fa9c835e
https://cdn-avatars.huggingface.co/v1/production/uploads/643839d9581e6bf0fa9c835e/JxlgR-zQhms-rfF0sDxD8.jpeg
Junda Chen
user
GindaChen
68f90580b9b2e4ae0467365a
Bo Pang
claimed_verified
2025-10-24T01:07:30.195Z
679ea9f83d9e2d122841ff06
/avatars/bfc157c6e824b140d57b83c35c34d453.svg
Bo Pang
user
PbPb314314
68f90580b9b2e4ae0467365b
Yi Gu
claimed_verified
2025-10-23T14:58:35.756Z
64c2ab94e818eec6128bfeb7
/avatars/459794e67d43a4b739ddcf90d1921b31.svg
Yi Gu
user
wuqing157
68f90580b9b2e4ae0467365c
Yibo Zhu
68f90580b9b2e4ae0467365d
Yimin Jiang
68f90580b9b2e4ae0467365e
Ion Stoica
68f90580b9b2e4ae0467365f
Eric Xing
68f90580b9b2e4ae04673660
Hao Zhang
Efficient Long-context Language Model Training by Core Attention
  Disaggregation
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.
CAD, a technique for long-context large language model training, improves throughput and balance by decoupling and distributing core attention computations.
core attention disaggregation core attention softmax(QK^T)V data parallel pipeline parallel token-level shards attention servers DistCA ping-pong execution in-place execution