Effizientes Training von Sprachmodellen mit langem Kontext durch Disaggregation der Kern-Attention

Efficient Long-context Language Model Training by Core Attention Disaggregation

October 20, 2025

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

cs.AI

papers.abstract

Wir präsentieren Core Attention Disaggregation (CAD), eine Technik, die das Training von Large Language Models mit langen Kontexten verbessert, indem die Kernberechnung der Attention, softmax(QK^T)V, vom Rest des Modells entkoppelt und auf einem separaten Pool von Geräten ausgeführt wird. In bestehenden Systemen ist die Kern-Attention mit anderen Schichten zusammengefasst; bei langen Kontextlängen führt ihr quadratischer Rechenaufwand im Vergleich zum nahezu linearen Wachstum anderer Komponenten zu Lastungleichgewichten und Verzögerungen in Daten- und Pipeline-Parallelgruppen. CAD wird durch zwei Beobachtungen ermöglicht. Erstens ist die Kern-Attention zustandslos: Sie hat keine trainierbaren Parameter und nur minimale temporäre Daten, sodass die Lastverteilung auf die Planung rechenintensiver Aufgaben reduziert wird. Zweitens ist sie komponierbar: Moderne Attention-Kernel behalten eine hohe Effizienz bei der Verarbeitung von fusionierten Batches von Token-Level-Shards mit beliebigen Längen. CAD unterteilt die Kern-Attention in Token-Level-Aufgaben und verteilt sie auf dedizierte Attention-Server, die Aufgaben dynamisch neu bündeln, um die Rechenlast auszugleichen, ohne die Kernel-Effizienz zu beeinträchtigen. Wir implementieren CAD in einem System namens DistCA, das ein Ping-Pong-Ausführungsschema verwendet, um die Kommunikation vollständig mit der Berechnung zu überlappen, und In-Place-Ausführung auf Attention-Servern, um den Speicherverbrauch zu reduzieren. Auf 512 H200 GPUs und Kontextlängen von bis zu 512k Tokens verbessert DistCA den end-to-end Trainingsdurchsatz um bis zu 1,35x, eliminiert Verzögerungen in Daten- und Pipeline-Parallelgruppen und erreicht eine nahezu perfekte Rechen- und Speicherbalance.

English

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.