TPLA: Tensor Parallel Latent Attention for Efficient Disaggregated
  Prefill \& Decode Inference
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.
68abc7c886b21a0e2e358a78
Xiaojuan Tang
claimed_verified
2025-08-25T08:19:43.131Z
6433c6ad7b8247480106c189
/avatars/42d84a7bbcfb9f5cc60bb460e6375f14.svg
Aurora
user
xiaojuan0920
68abc7c886b21a0e2e358a79
Fanxu Meng
claimed_verified
2025-08-26T09:50:36.462Z
643f55d4ec817b766686438a
/avatars/0feb460432c92ab9ada0d417a7a38f6a.svg
mengfanxu
user
fxmeng
68abc7c886b21a0e2e358a7a
Pingzhi Tang
68abc7c886b21a0e2e358a7b
Yuxuan Wang
68abc7c886b21a0e2e358a7c
Di Yin
68abc7c886b21a0e2e358a7d
Xing Sun
claimed_verified
2025-09-15T15:09:35.760Z
647401e50da364bd0d002f2a
https://cdn-avatars.huggingface.co/v1/production/uploads/no-auth/vPuPn7EV092mLBOM2YZXd.png
XING SUN
user
tedsun
68abc7c886b21a0e2e358a7e
Muhan Zhang
TPLA: Tensor Parallel Latent Attention for Efficient Disaggregated
  Prefill \& Decode Inference
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.
Tensor-Parallel Latent Attention (TPLA) enhances tensor parallelism efficiency by partitioning latent representations and input dimensions, preserving the benefits of compressed key-value caches while maintaining strong representational capacity.
Multi-Head Latent Attention (MLA) tensor parallelism (TP) latent representation attention heads Grouped Query Attention (GQA) Tensor-Parallel Latent Attention (TPLA) Grouped Latent Attention (GLA) prefilling tensor-parallel decoding Hadamard transform PCA FlashAttention-3 commonsense LongBench benchmarks