TPLA: Attenzione Latente con Parallelismo Tensoriale per un'Inferenza Efficiente di Prefill e Decodifica Disaggregata

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

August 21, 2025

Autori: Xiaojuan Tang, Fanxu Meng, Pingzhi Tang, Yuxuan Wang, Di Yin, Xing Sun, Muhan Zhang

cs.AI

Abstract

Multi-Head Latent Attention (MLA), introdotta in DeepSeek-V2, comprime gli stati chiave-valore in un vettore latente a basso rango, memorizzando solo questo vettore per ridurre l'uso di memoria. Tuttavia, nel parallelismo tensoriale (TP), le teste di attenzione vengono calcolate su più dispositivi, e ciascun dispositivo deve caricare l'intera cache, riducendo il vantaggio di MLA rispetto a Grouped Query Attention (GQA). Proponiamo Tensor-Parallel Latent Attention (TPLA): uno schema che partiziona sia la rappresentazione latente che la dimensione di input di ciascuna testa tra i dispositivi, esegue l'attenzione in modo indipendente per ciascuna partizione e poi combina i risultati con un'operazione all-reduce. TPLA preserva i benefici di una cache KV compressa sbloccando l'efficienza del TP. A differenza di Grouped Latent Attention (GLA), ogni testa in TPLA continua a sfruttare l'intera rappresentazione latente, mantenendo una capacità rappresentativa più forte. TPLA è compatibile con modelli pre-addestrati utilizzando MLA: supporta il prefilling in stile MLA e abilita il decoding efficiente in parallelo tensoriale senza bisogno di riaddestramento. Applicando semplici trasformazioni ortogonali -- ad esempio, la trasformata di Hadamard o PCA -- prima del partizionamento TP si mitiga ulteriormente l'interferenza tra partizioni, ottenendo un degrado minimo dell'accuratezza. Riducendo la cache KV per dispositivo in DeepSeek-V3 e Kimi-K2, otteniamo rispettivamente accelerazioni di 1.79x e 1.93x, con una lunghezza di contesto di 32K token, mantenendo le prestazioni sui benchmark di commonsense e LongBench. TPLA può essere implementata con FlashAttention-3, consentendo un'accelerazione pratica end-to-end.

English

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.