Track papers across labs, topics, and trends.
Search papers, labs, and topics across Lattice.
This paper investigates the limitations of the Multi-head Latent Attention (MLA) implementation in Megatron-Core, specifically the restriction against using its absorbed form during training due to significant memory overhead. The authors quantify this overhead, showing that using the absorbed form inflates activation memory by 20-34%, which can lead to device-fit issues at scale. To address this, they introduce LAGA (Latent All-Gather Attention), which maintains the communication efficiency of the absorbed form while avoiding its memory pitfalls, achieving a 1.98x reduction in collective communication without compromising memory performance.
LAGA achieves nearly identical performance to the explicit MLA implementation while slashing collective communication by nearly 2x, challenging the notion that efficiency must come at the cost of memory overhead.
Multi-head Latent Attention (MLA) ships two implementations in Megatron-Core: an explicit form used for training and an absorbed form -- which slashes collective communication by gathering only the compressed latent -- that is fully implemented but hard-asserted out of training (the forward opens with"assert not (self.training and self.cache_mla_latents)"), allowed only in inference decode. The library documents no reason. We show the restriction is well-founded and quantify why: ported to training, the absorbed form is a memory trap -- its intermediates live in n_h x d_kv dimensions per token, larger than the per-head K/V they replace -- inflating activation memory by 20-34%, up to 9.2 GB at DeepSeek-V3 scale (n_h=128, seq=16384, SP=8, eager kernel; the gap widens to 19.2 GB under a fused kernel), enough to change device-fit. This measurement, validated on two axes (linear in seq and n_h) and cross-verified on NVIDIA A100, explains the otherwise-undocumented restriction and leaves practitioners with no low-communication MLA training path. We then provide one. LAGA (Latent All-Gather Attention) keeps the absorbed form's latent-gather communication but rejects the absorb reformulation, instead reconstructing per-head K/V locally from the gathered latent. On 8x Ascend 910B at real DeepSeek-V3 dimensions, LAGA cuts collective communication 1.98x, matches explicit memory within 0.5%, is bit-identical to explicit at SP=1 and equivalent to within 1e-3 at SP=2-8, and under a fused attention kernel improves attention-block throughput 1.04-1.06x single-node and 1.07-1.24x cross-node -- leading at all sequence lengths in the cross-node regime MLA is deployed for.