TPLA: Тензорно-параллельное латентное внимание для эффективного распределенного вывода на этапах предзаполнения и декодирования

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

August 21, 2025

Авторы: Xiaojuan Tang, Fanxu Meng, Pingzhi Tang, Yuxuan Wang, Di Yin, Xing Sun, Muhan Zhang

cs.AI

Аннотация

Многоголовое латентное внимание (Multi-Head Latent Attention, MLA), представленное в DeepSeek-V2, сжимает ключевые и значения в низкоранговый латентный вектор, кэшируя только этот вектор для уменьшения использования памяти. Однако при тензорном параллелизме (Tensor Parallelism, TP) внимание вычисляется на нескольких устройствах, и каждое устройство должно загружать полный кэш, что снижает преимущество MLA перед групповым запросным вниманием (Grouped Query Attention, GQA). Мы предлагаем Тензорно-Параллельное Латентное Внимание (Tensor-Parallel Latent Attention, TPLA): схему, которая разделяет как латентное представление, так и входное измерение каждой головы между устройствами, выполняет внимание независимо для каждого сегмента, а затем объединяет результаты с помощью операции all-reduce. TPLA сохраняет преимущества сжатого кэша ключей и значений (KV cache), одновременно обеспечивая эффективность TP. В отличие от Группового Латентного Внимания (Grouped Latent Attention, GLA), каждая голова в TPLA по-прежнему использует полное латентное представление, сохраняя более высокую репрезентативную способность. TPLA полностью совместима с моделями, предварительно обученными с использованием MLA: она поддерживает предварительное заполнение в стиле MLA и обеспечивает эффективное тензорно-параллельное декодирование без необходимости переобучения. Применение простых ортогональных преобразований — например, преобразования Адамара или PCA — перед разделением на сегменты TP дополнительно снижает межсегментные помехи, что приводит к минимальной потере точности. Уменьшая размер кэша KV на устройстве для DeepSeek-V3 и Kimi-K2, мы достигаем ускорения в 1.79x и 1.93x соответственно при длине контекста в 32K токенов, сохраняя производительность на тестах commonsense и LongBench. TPLA может быть реализована с использованием FlashAttention-3, что обеспечивает практическое ускорение на всех этапах.

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.