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

Xiaojuan Tang, Fanxu Meng, Pingzhi Tang, Yuxuan Wang, Di Yin, Xing Sun and 1 other

DeepSeek's Multi-Head Latent Attention shrinks the KV cache by compressing it into one shared latent vector — but that sharing backfires when you spread the model across several GPUs, because every GPU must then keep the whole latent vector, wiping out the savings; this paper (Peking University, Tencent — Xiaojuan Tang, Fanxu Meng, Muhan Zhang, and colleagues) presents TPLA, which slices the latent vector across GPUs yet lets every attention head still see all of it (via an all-reduce), converts existing DeepSeek checkpoints without retraining using Hadamard or PCA transforms, and uses plain latent attention for prefill and the sliced version for decode — for up to 1.93 times faster decoding at 32K tokens with accuracy preserved.

The argument in one glance

Problem

The trick that shrinks the KV cache breaks when the model is split across GPUs

Generating text with a large language model is memory-bound: each new token requires reading the key-value (KV) cache of all previous tokens, so the size of that cache, and the bandwidth to read it, set the speed. Multi-Head Latent Attention (MLA), introduced in DeepSeek-V2, is a powerful way to shrink it: instead of caching separate key and value vectors for every attention head, it compresses them into a single small low-rank latent vector, and caches only that. This gives a much smaller KV cache than the usual alternative, Grouped Query Attention (GQA), while keeping more of the model's quality, which is why DeepSeek models use it. But large models do not run on one GPU; they are split across several using tensor parallelism, where each GPU computes a subset of the attention heads. Here MLA hits a wall. Because the compressed latent vector is shared by all the heads, every GPU that computes some heads must load the entire latent vector — it cannot use just a slice of it. So under tensor parallelism the latent cache is replicated in full on every device, and MLA's per-device memory advantage over GQA evaporates. In DeepSeek-V3, for example, the latent KV dimension is 576 per token, and that full 576 must sit on every GPU no matter how many GPUs share the work — often a larger per-device footprint than an equivalent GQA model would have.
What people did

Split the heads into groups, each with its own slice of the latent

The known fix is Grouped Latent Attention (GLA). It divides both the attention heads and the latent representation into a few groups — typically two — so that each group of heads loads only its own portion of the latent vector rather than the whole thing. Each GPU then holds just its group's slice of the latent cache, restoring per-device memory savings under tensor parallelism. This directly attacks the replication problem by making the latent, like the heads, something that can be partitioned across devices.
Why it fell short

Grouping shrinks each head's view and forces training from scratch

Grouped Latent Attention pays for its memory savings in two ways. First, representational capacity: by assigning each group of heads only its own slice of the latent, GLA means each attention head no longer attends to the full latent representation but only to a fraction of it. Since MLA's quality advantage came precisely from every head drawing on the rich shared latent, cutting each head down to a slice weakens what the heads can express — the memory win is bought with a modeling loss. Second, deployability: GLA is a different architecture, so a model must be trained from scratch to use it, which costs enormous compute and cannot reuse the many strong models already pre-trained with MLA, like DeepSeek. So GLA fixes the tensor-parallelism footprint but degrades the attention and abandons existing checkpoints. What was missing was a way to give each device only a slice of the latent cache — the memory win — while still letting every head attend to the full latent representation, and to do it on top of an already-trained MLA model rather than requiring a new one.
This paper's idea

Slice the latent across GPUs but let every head still attend to all of it

TPLA (Tensor Parallel Latent Attention) partitions the latent representation across devices and also splits each attention head's input dimension across them, so each GPU stores only a slice of the latent cache. The crucial difference from GLA is what each head sees: in TPLA every head still attends to the full latent space, because each device computes its head's attention over its own latent slice, and the partial results are then combined across devices with an all-reduce. So the head's final output reflects the entire latent, preserving MLA's representational capacity, while no single device ever holds more than a partition of the cache — the per-device memory win without the modeling loss. Three further pieces make this practical. First, to slice a shared latent vector cleanly across devices without the softmax and normalization introducing cross-shard discrepancies, TPLA applies orthogonal reparameterizations of RMSNorm and softmax — realized with Hadamard or PCA transforms — before partitioning, which spread information evenly across the latent dimensions so each slice carries a fair share and accuracy barely drops. Second, because of this, TPLA is drop-in compatible with existing MLA checkpoints: a pre-trained DeepSeek model can be converted to TPLA with only a minor, easily-recovered performance drop, no training from scratch. Third, since prefill is compute-bound while decode is memory-bound, TPLA keeps the exact MLA form during prefill (where memory is not the constraint and accuracy is best preserved) and switches to TPLA only for decode (where the smaller per-device cache speeds things up), minimizing conversion-induced error. Reducing the per-device KV cache for DeepSeek-V3 and Kimi-K2, TPLA achieves 1.79 and 1.93 times decoding speedups at a 32K-token context while maintaining accuracy on commonsense and LongBench benchmarks, and it can be implemented on top of FlashAttention-3.

Why this problem is worth solving

The KV cache is the dominant cost of long-context LLM serving, so any technique that shrinks it is valuable — but only if it survives the way models are actually deployed.

When a model generates text, every new token attends over the stored keys and values of all previous tokens, so the KV cache grows with the context length and must be read in full at each step. For long contexts this cache dominates memory use and, because decoding is memory-bandwidth-bound, dominates speed too — reading a smaller cache directly means faster tokens. MLA is prized because it compresses this cache dramatically while keeping quality high, which is a large win at exactly the point where serving is most expensive. But a compression technique is only useful if it composes with tensor parallelism, because the models worth serving are too large for one GPU and must be split across many. If a compression scheme forces every device to hold the whole compressed cache, then adding GPUs does not reduce the per-device memory pressure that limits how long a context or how large a batch each GPU can handle — the compression helps the total but not the per-device footprint that actually gates deployment. So the value here is in making MLA’s excellent compression actually pay off under the tensor parallelism real deployments use, turning a technique that looked good on a single GPU into one that scales. Getting the per-device cache down is what lets a fixed cluster serve longer contexts and more users, which at the scale of modern LLM services is a direct cost saving.

The deeper insight is that the shared latent and per-device partitioning were assumed to be at odds, and that a small algebraic move — combining sliced attention with an all-reduce and an orthogonal transform — reconciles them.

It seems like a genuine dilemma: MLA’s power comes from every head reading one shared latent vector, but tensor parallelism needs that vector partitioned so each device holds only a piece. GLA resolves it by giving each head only a piece, accepting the capacity loss. TPLA’s realization is that attention over a full vector can be computed as a sum of attentions over its slices, so you can distribute the slices across devices, have each device compute the partial attention over its slice, and add the pieces back together with an all-reduce — recovering exactly the result a head would have gotten from the whole latent, but with each device only ever storing and reading its slice. The head’s view is full; only the storage is partitioned. The subtlety is that operations like softmax and layer normalization do not split cleanly across arbitrary slices, so a naive partition introduces cross-shard error. The fix is to first apply an orthogonal transform — a Hadamard rotation or a PCA rotation — that redistributes the latent’s information evenly across its dimensions, so that any slice carries a representative share and the recombination is accurate. Because these transforms are orthogonal, they are mathematically invertible and can be folded into an existing model’s weights, which is why an already-trained MLA checkpoint can be converted without retraining. Recognizing that partitioning storage need not mean partitioning the head’s view, and that an orthogonal reparameterization makes the slicing clean, is what turns a seeming conflict into a drop-in speedup.

What people did before, and where it broke

Multi-Head Latent Attention (MLA) under tensor parallelism. Compress the KV cache into one shared latent vector and split the heads across GPUs. Where it broke. The shared latent must be replicated in full on every device, so the per-device KV cache does not shrink with more GPUs, erasing MLA’s advantage over GQA.
Grouped Query Attention (GQA). Share KV across groups of query heads, which partitions cleanly under tensor parallelism. Where it broke. Its representational capacity is inferior to MLA, so quality suffers relative to latent attention.
Grouped Latent Attention (GLA). Split heads and latent into groups so each group loads only its slice. Where it broke. Each head then attends to only a fraction of the latent, weakening representation, and it is a new architecture that must be trained from scratch, abandoning existing MLA checkpoints.
Generic KV-cache compression (pruning, merging, quantization). Reduce the cache by discarding or approximating entries. Where it broke. These alter or drop KV information and can degrade quality, and they do not address the tensor-parallel replication problem specific to MLA’s shared latent.
The gap. No method gave each device only a slice of the latent cache while keeping every head’s attention over the full latent, deployable on an already-trained MLA model without retraining.

How TPLA works

From a replicated shared latent to a sliced latent with full-view attention

↓ →
new
↓ →
new
↓ →
new
↓ →
new
↓ →
new

The idea

Shrink the KV cache to speed up memory-bound decoding. TPLA’s purpose is to reduce the per-device KV cache that bottlenecks long-context generation, so each token is produced faster. Cutting the memory the key-value cache consumes so that memory-bound decoding runs faster is the same KV-reduction goal the work on efficient long-context reasoning with speculative context sparsity (SpeContext) covered elsewhere on this site pursues, attacking the cache that dominates long-context inference. TPLA attacks it specifically for latent attention under tensor parallelism, where the compressed cache would otherwise be replicated.

Partition the computation across devices to fit and scale. TPLA splits the latent and each head across GPUs and recombines with an all-reduce, a tensor-parallel decomposition of the attention. Distributing a model’s tensors and computation across multiple devices so that large models fit and run in parallel is the same tensor-parallel-partitioning approach the work on composing SPMD partitioning strategies for machine learning (PartIR) covered elsewhere on this site is built on, sharding computation across a device mesh. TPLA is a new sharding for latent attention that, unlike the naive one, avoids replicating the shared cache.

Transform the representation so aggressive compression stays lossless. TPLA rotates the latent with a Hadamard or PCA transform so that slicing it across devices barely loses accuracy, folding the transform into the weights. Applying an orthogonal transform to redistribute a tensor’s information so it survives an aggressive size reduction with minimal error is the same encode-for-compression idea the work on a compact encoding that packs tensor values tightly (Atalanta) covered elsewhere on this site embodies, reshaping data so it compresses cleanly. TPLA uses invertible rotations to make the shared latent evenly sliceable, which is what lets it convert an existing MLA model without retraining.

Following a decode step through TPLA

From a converted MLA checkpoint to a fast tensor-parallel decode

A trace, drawn from the paper, of how TPLA runs one decoding step across GPUs after converting an MLA model.

Step 1 of 5

Convert the pre-trained MLA model offline

Starting from a DeepSeek MLA checkpoint, TPLA applies an orthogonal (Hadamard or PCA) reparameterization to the latent-related weights and folds it in. Because the transform is invertible, the model computes the same function, but its latent is now evenly sliceable.

This conversion needs no training from scratch, only minor recovery.

Why it works

Attention over a whole vector equals the sum over its slices. A shared latent seems to force full replication → but computing partial attention per slice and adding via all-reduce → so each head’s full-latent result is recovered while each device holds only a slice.
Slicing storage need not slice the head’s view. GLA shrank memory by shrinking each head’s view → but TPLA partitions only the storage, not the attention → so per-device memory falls without losing representational capacity.
An orthogonal transform makes the slices fair. Softmax and normalization do not split cleanly across arbitrary slices → but a Hadamard or PCA rotation spreads information evenly first → so every slice carries a representative share and accuracy barely drops.
Invertible transforms fold into existing weights. A new attention scheme would need training from scratch → but orthogonal reparameterization preserves the model’s function → so a pre-trained MLA checkpoint converts with only minor, recoverable loss.
Splitting prefill and decode limits conversion error. Converting everywhere would add error where it is not needed → but keeping exact MLA for compute-bound prefill and TPLA only for memory-bound decode → so the speedup lands in decode while accuracy is preserved in prefill.

What it costs (trade-offs and scope)

Conversion still costs a little accuracy. Reparameterizing and slicing an MLA model incurs a minor performance drop; the paper reports it is easily recovered, but it is not exactly free, which is partly why prefill stays in exact MLA form.
The all-reduce adds communication. Recombining partial attention outputs across devices requires an all-reduce each step, so TPLA trades some inter-GPU communication for the per-device memory savings — a good trade when decoding is memory-bound, but a real cost.
The benefit is specific to MLA under tensor parallelism. TPLA targets the replication problem of latent attention when split across GPUs; on a single device, or for non-latent attention, its particular advantage does not apply.
Prefill and decode use different forms. Running MLA for prefill and TPLA for decode is a disaggregated design that a serving system must support switching between, adding implementation complexity compared with a single uniform attention.
Results are on specific models, contexts, and benchmarks. The 1.79 and 1.93 times speedups are for DeepSeek-V3 and Kimi-K2 at 32K-token context on commonsense and LongBench evaluations; other models, context lengths, or tasks may differ.

The comparison

Attention schemes under tensor parallelism

Property	MLA	GLA	TPLA
Per-device KV cache under TP
What each head attends to
Needs training from scratch
Clean slicing
Reported decode speedup

A comparison of latent-attention schemes under tensor parallelism. Plain MLA replicates the full shared latent on every device, losing its memory advantage; GLA slices the latent per group but shrinks each head's view and requires training from scratch. TPLA slices the latent per device yet keeps every head attending to the full latent via an all-reduce, uses Hadamard or PCA transforms to slice cleanly, and converts pre-trained MLA checkpoints without retraining — with MLA used for compute-bound prefill and TPLA for memory-bound decode, giving up to 1.93 times faster decoding at 32K tokens with accuracy preserved. Exact figures should be read from the paper.Underlined values explain themselves when clicked.

Where this stops working

What to keep in mind

This explainer is based on the paper’s full text (the authors’ arXiv copy, 2508.15881, of the ASPLOS 2026 paper). The described method is the paper’s: Multi-Head Latent Attention (MLA, from DeepSeek-V2) caches a single shared low-rank latent vector to shrink the KV cache, but under tensor parallelism every device must replicate the full latent, erasing MLA’s per-device advantage over Grouped Query Attention. TPLA partitions the latent and each head’s input dimension across devices, computes attention per slice, and aggregates with an all-reduce, so each device holds only a slice while every head still attends to the full latent — unlike Grouped Latent Attention (GLA), which reduces each head’s view and requires training from scratch. TPLA applies orthogonal reparameterizations of RMSNorm and softmax (via Hadamard or PCA transforms) before slicing to mitigate cross-shard discrepancies, making it drop-in compatible with pre-trained MLA checkpoints with only minor, recoverable loss; and it keeps MLA for compute-bound prefill and TPLA for memory-bound decode to minimize conversion error. The reported results — 1.79 and 1.93 times decoding speedups for DeepSeek-V3 and Kimi-K2 at a 32K-token context, with accuracy maintained on commonsense and LongBench benchmarks, and implementability on FlashAttention-3 — are from the paper’s evaluation. Background on the KV cache, MLA, GQA, GLA, tensor parallelism, softmax, RMSNorm, and orthogonal transforms is standard context. The step-by-step walkthrough illustrates the described mechanism. Exact derivations, the reparameterization details, and full results should be read from the paper.

Background

Why efficient LLM inference keeps redesigning attention around the KV cache

It helps to see TPLA within the sustained effort to make LLM inference cheap, most of which circles back to the key-value cache. Because generation is memory-bound and the cache grows with context length, the cache is the central object that serving systems and attention designers keep trying to shrink: by compressing it (as MLA does with a low-rank latent), by pruning or quantizing it, by sharing it across query heads (as GQA does), or by managing its memory cleverly at serving time. Each such idea then has to survive contact with how models are actually run — across many GPUs, in distinct prefill and decode phases, on top of optimized attention kernels — and techniques that look good in isolation often stumble on those realities, as MLA does on tensor parallelism. The recurring pattern is a co-design between the attention math and the systems it runs on, so that a compression that saves memory also partitions across devices and slots into existing kernels and checkpoints. This drive to serve LLMs efficiently by carefully engineering the phases and mechanics of inference is the same terrain the work on high-goodput serving via prefill-decode multiplexing (MuxWise) covered elsewhere on this site occupies, tuning how the compute- and memory-bound phases share the hardware. TPLA’s distinctive contribution is to fix latent attention’s tensor-parallel flaw with an algebraic decomposition plus an orthogonal reparameterization, so an already-trained DeepSeek model gets its memory savings back under real multi-GPU serving.

If you remember only three things

1. LLM decoding is memory-bound: each new token must read the key-value (KV) cache of all prior tokens, so shrinking that cache speeds generation. Multi-Head Latent Attention (MLA), from DeepSeek-V2, compresses the KV states into a single small shared low-rank latent vector and caches only that, beating Grouped Query Attention on both memory and quality. But large models run split across GPUs via tensor parallelism, where each GPU computes some attention heads — and because MLA's latent is shared by all heads, every GPU must replicate the whole latent vector. So under tensor parallelism MLA's per-device cache does not shrink, erasing its advantage. The prior fix, Grouped Latent Attention (GLA), slices the latent per group of heads to save per-device memory, but each head then sees only its slice (losing representational capacity), and GLA is a new architecture that must be trained from scratch, abandoning existing MLA checkpoints.
2. TPLA partitions the latent representation and each attention head across devices so each GPU stores only a slice of the latent cache, but — crucially — lets every head still attend to the full latent: each device computes its head's partial attention over its slice, and the partial results are combined with an all-reduce, recovering exactly the full-latent result. So the per-device memory shrinks without the capacity loss GLA suffers. To make slicing clean despite softmax and normalization not splitting across arbitrary slices, TPLA first applies orthogonal reparameterizations (Hadamard or PCA transforms) that redistribute the latent's information evenly; because these are invertible, they fold into the weights, so a pre-trained MLA model converts to TPLA without retraining and with only minor, recoverable loss. And since prefill is compute-bound while decode is memory-bound, TPLA keeps exact MLA for prefill and uses TPLA only for decode, minimizing conversion-induced error.
3. The key realization is that partitioning storage need not partition the head's view: attention over the whole latent equals the sum of attentions over its slices, so distributing the slices with an all-reduce gives memory savings and full capacity at once, and an orthogonal transform makes the slicing accurate on an existing checkpoint. Reducing the per-device KV cache for DeepSeek-V3 and Kimi-K2, TPLA achieves 1.79 and 1.93 times decoding speedups at a 32K-token context while maintaining accuracy on commonsense and LongBench benchmarks, and it can be implemented on top of FlashAttention-3. The trade-offs: conversion costs a little (recoverable) accuracy, the all-reduce adds inter-GPU communication, the benefit is specific to MLA under tensor parallelism, and running different attention forms for prefill and decode adds serving complexity.

Check yourself

Check your understanding

1. Why does Multi-Head Latent Attention (MLA) lose its advantage under tensor parallelism?
2. What is the limitation of Grouped Latent Attention (GLA), the prior fix?
3. How does TPLA slice the latent across devices yet keep each head attending to all of it?
4. How does TPLA convert existing MLA models without retraining, and split prefill from decode?

Sources

The paper is “TPLA: Tensor Parallel Latent Attention for Efficient Disaggregated Prefill and Decode Inference” by Xiaojuan Tang, Fanxu Meng, Pingzhi Tang, Yuxuan Wang, Di Yin, Xing Sun, and Muhan Zhang (Peking University and Tencent, among the affiliations), published at ASPLOS 2026, DOI 10.1145/3779212.3790237. It is available on arXiv as 2508.15881.

This explainer is written from the paper’s full text (the authors’ arXiv copy). Multi-Head Latent Attention (MLA) compresses the KV cache into a shared low-rank latent vector, but under tensor parallelism every device must replicate the full latent, erasing its advantage over Grouped Query Attention; Grouped Latent Attention (GLA) slices the latent per group but reduces each head’s view and needs training from scratch. TPLA partitions the latent and each head across devices, computes partial attention per slice, and aggregates with an all-reduce, so each device holds only a slice while every head still attends to the full latent; it applies orthogonal (Hadamard or PCA) reparameterizations before slicing to keep accuracy, making it drop-in compatible with pre-trained MLA checkpoints, and it uses MLA for compute-bound prefill and TPLA for memory-bound decode. The paper reports 1.79 and 1.93 times decoding speedups for DeepSeek-V3 and Kimi-K2 at a 32K-token context, with accuracy maintained on commonsense and LongBench benchmarks, and implementability on FlashAttention-3. Background on KV caching, MLA, GQA, GLA, tensor parallelism, and orthogonal transforms is standard context. The numbers are the paper’s own; exact mechanisms and full results should be confirmed in the paper.

TPLA: Tensor Parallel Latent Attention for Efficient Disaggregated Prefill &#38; Decode Inference