TPLA: Tensor Parallel Latent Attention for Efficient Disaggregated Prefill and Decode Inference

TL;DR

TPLA introduces a tensor parallel attention scheme that maintains the benefits of latent key-value caching while enabling efficient distributed inference, resulting in significant speedups with minimal accuracy loss.

Contribution

It proposes Tensor-Parallel Latent Attention (TPLA), a novel method that combines latent caching with tensor parallelism, compatible with pre-trained models and supporting efficient decoding.

Findings

Achieves 1.79x and 1.93x speedups on DeepSeek-V3 and Kimi-K2.

Maintains performance on benchmarks with minimal accuracy degradation.

Supports integration with FlashAttention-3 for practical acceleration.

Abstract

Multi-Head Latent Attention (MLA), introduced in DeepSeek-V2, compresses key-value states into a low-rank latent vector, 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 cache, eroding the advantage of MLA over Grouped Query Attention (GQA). We propose Tensor-Parallel Latent Attention (TPLA): a scheme that partitions both the latent representation and each head's input dimension across devices, performs attention independently per shard, and then combines results with an all-reduce. TPLA preserves the benefits of a compressed KV cache while unlocking TP efficiency. Unlike Grouped Latent Attention (GLA), every head in TPLA still leverages the full latent representation, maintaining stronger representational capacity. TPLA is drop-in compatible with models pre-trained using MLA: it supports MLA-style prefilling and enables efficient tensor-parallel decoding without retraining. Applying simple orthogonal transforms -- e.g., the Hadamard transform or PCA -- before TP slicing further mitigates cross-shard interference, yielding minimal accuracy degradation. By reducing the per-device KV cache for DeepSeek-V3 and Kimi-K2, we achieve 1.79x and 1.93x speedups, respectively, at a 32K-token context length while maintaining performance on commonsense and LongBench benchmarks. TPLA can be implemented with FlashAttention-3, enabling practical end-to-end acceleration.