一句话：提出 TPLA，按设备切分 MLA 的潜在表示与头输入维度，在张量并行下降低单设备 KV 缓存并保持精度与加速推理。
MLA 只缓存低秩潜在向量以压缩 KV 显存，但在张量并行下每台设备仍需加载完整 cKV，削弱了其相对 GQA 的优势。本文提出 TPLA：同时切分潜在表示与每个注意力头的输入维度，各分片独立计算注意力后 all-reduce 聚合；每个头仍关注完整潜在空间，从而保留 MLA 表达能力并降低单设备缓存。为兼容 MLA 检查点，作者用 Hadamard 与 PCA 变换对 RMSNorm 与 softmax 做正交重参数化以缓解跨分片偏差，并让 compute-bound 的 prefill 保持 MLA、memory-bound 的 decode 切换为 TPLA。在 DeepSeek-V3 与 Kimi-K2（32K 上下文）上分别取得 1.79× 与 1.93× 加速，精度保持，并可基于 FlashAttention-3 实现。
Multi-Head Latent Attention (MLA), introduced in DeepSeek-V2, compresses key–value states into a low-rank latent vector cKV, caching only this vector to reduce memory. In tensor parallelism (TP), however, attention heads are computed across multiple devices, and each device must load the full cKV, eroding the advantage of MLA over Grouped Query Attention (GQA). We present TPLA, a scheme that partitions both the latent representation and each head's input dimension across devices, performs attention independently on each shard, and aggregates the results with an all-reduce. Unlike GLA, every attention head in TPLA still attends to the full latent space, preserving MLA's representational capacity while reducing the per-device KV cache. To make TPLA drop-in compatible with MLA checkpoints, we further derive orthogonal reparameterizations of RMSNorm and softmax---instantiated with Hadamard and PCA transforms---that mitigate cross-shard discrepancies when slicing latent vectors across devices. Finally, we introduce a prefill-decode separation scheme that keeps the MLA form during compute-bound prefilling and switches to TPLA during memory-bound decoding, minimizing conversion-induced error. By reducing the per-device KV cache for DeepSeek-V3 and Kimi-K2, we achieve 1.79x) and 1.93) speedups respectively, at a 32K-token context length while maintaining accuracy on commonsense and LongBench benchmarks. TPLA can be further implemented on top of FlashAttention-3, enabling practical end-to-end acceleration.
@inproceedings{TangMTWY0Z26,
title = {TPLA: Tensor Parallel Latent Attention for Efficient Disaggregated Prefill & Decode Inference},
author = {Xiaojuan Tang and Fanxu Meng and Pingzhi Tang and Yuxuan Wang and Di Yin and Xing Sun and Muhan Zhang},
booktitle = {ASPLOS},
year = {2026},
doi = {10.1145/3779212.3790237}
}