TPLA: Tensor-Parallel Latent Attention beschleunigt Nachlauf großer Modelle

Die neue Methode Tensor-Parallel Latent Attention (TPLA) kombiniert die Speicher‑Effizienz der Multi‑Head Latent Attention (MLA) mit der Rechenleistung von Tensor‑Parallelism. MLA reduziert die Größe der Key‑Value‑Cache‑Daten, indem sie diese in einen kompakten latenten Vektor zusammenfasst. In herkömmlichem Tensor‑Parallel‑Setup muss jedes Gerät jedoch den gesamten Cache laden, was den Speicher‑Vorteil von MLA zunichte macht.

Mit TPLA wird der latente Vektor sowie die Eingabedimension jeder Attention‑Head auf die verschiedenen Geräte aufgeteilt. Jeder Shard führt die Attention selbstständig durch und die Ergebnisse werden anschließend über einen All‑Reduce‑Befehl zusammengeführt. So bleibt die Speicher‑Kompression erhalten, während die Parallelität voll ausgenutzt wird.

Im Gegensatz zu Grouped Latent Attention (GLA) nutzt jede Head in TPLA weiterhin die komplette latente Darstellung, was die Repräsentationskraft erhöht. TPLA ist vollständig kompatibel mit bereits vortrainierten Modellen, die MLA verwendet haben. Es ermöglicht MLA‑basierte Prefilling‑ und Decoding‑Schritte ohne erneutes Training.

Durch die Anwendung einfacher orthogonaler Transformationen – etwa der Hadamard‑ oder PCA‑Transform – vor der Aufteilung in Shards kann die Interferenz zwischen den Shards weiter reduziert werden, ohne die Genauigkeit wesentlich zu beeinträchtigen.

In Tests mit DeepSeek‑V3 und Kimi‑K2 erzielte TPLA bei einer 32‑K‑Token‑Kontextlänge Geschwindigkeitssteigerungen von 1,79‑ und 1,93‑fach, während die Leistung auf Commonsense‑ und LongBench‑Benchmarks erhalten blieb. Die Implementierung lässt sich nahtlos mit FlashAttention‑3 verbinden, was eine praktische End‑zu‑End‑Beschleunigung ermöglicht.