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

ArXiv: 2508.15881

🎯 Pitch

Multi-Head Latent Attention (MLA) saves memory by compressing the KV cache into one shared vector, but this advantage vanishes under tensor parallelism because every device must load that same full vector. TPLA solves this by slicing the compressed representation across devices so each shard only stores a fraction of it, yet every attention head still sees the complete information—yielding up to 1.93× decoding speedup at 32K context with negligible accuracy loss.

---

1. Executive Summary

This paper introduces Tensor Parallel Latent Attention (TPLA), a scheme that partitions both the latent KV representation and each attention head’s input dimension across devices, performs attention independently per shard, and combines results with an all-reduce—preserving the compressed KV cache benefits of Multi-Head Latent Attention while enabling efficient tensor parallelism without retraining. Evaluated on DeepSeek-V3 and Kimi-K2 across commonsense benchmarks and LongBench, TPLA uses two reparameterization strategies (Hadamard transform and PCA-based orthogonal transforms) to mitigate cross-shard interference in RMSNorm slicing and softmax slicing during the conversion from MLA, combined with a prefill–decode separation that retains MLA during compute-bound prefilling. TPLA achieves 1.79× and 1.93× decoding throughput speedups over MLA at a 32K-token context length for DeepSeek-V3 and Kimi-K2 respectively, with training-free conversion limiting accuracy degradation to under 2.15% on LongBench, establishing that the per-device KV cache memory bottleneck in tensor-parallel MLA inference can be substantially alleviated without sacrificing representational capacity only when the latent dimension is partitioned such that every attention head retains access to the full latent representation across devices.

2. Context and Motivation

The Core Problem: MLA's Memory Efficiency Dissolves Under Tensor Parallelism

The fundamental tension that TPLA addresses is a system-level clash between two independently valuable inference optimizations: Multi-Head Latent Attention (MLA) and tensor parallelism (TP). Understanding why this clash occurs — and why it matters at the scale of modern LLM deployment — requires stepping through the basic mechanics of both techniques.

When an LLM generates text autoregressively, every new token requires computing attention over all previous tokens. The intermediate key and value tensors for those previous tokens are stored in the KV cache to avoid recomputation. For a model with $h$ attention heads and head dimension $d_h$, the per-token KV cache size is $2 \times h \times d_h$ (keys and values, each with $h$ vectors of dimension $d_h$). As context lengths grow to 32K, 128K, or beyond, this cache becomes enormous — it is the dominant memory consumer during decoding and, because it must be read from high-bandwidth memory (HBM) at every generation step, it also becomes the primary memory bandwidth bottleneck.

MLA, introduced in DeepSeek-V2, tackles this directly. Rather than storing full-size keys and values per head, MLA compresses them into a single low-rank latent vector $\mathbf{c}^{\mathrm{KV}}$ of dimension $4d_h$ (512 for DeepSeek-V3's $d_h = 128$), plus a small RoPE positional component of dimension $d_r$ (64 for DeepSeek-V3). The per-token KV cache shrinks to $4d_h + d_r = 576$ dimensions, a dramatic reduction from the $2 \times h_{kv} \times d_h$ dimensions that standard multi-head or grouped-query attention would require. Critically, MLA achieves this compression during pretraining — the low-rank projection matrices $W^{DKV}$, $W^{UK}$, and $W^{UV}$ are learned parameters, not post-hoc approximations. This means MLA models are trained end-to-end with the compressed representation, and the compression is not lossy in the way that post-training pruning or quantization would be.

Tensor parallelism, on the other hand, is a distributed inference technique that splits individual layers across multiple GPUs. In the attention module, a common approach is to partition attention heads: with TP=2, GPU 0 computes attention for heads 0 through $h/2-1$, GPU 1 computes attention for heads $h/2$ through $h-1$, and the outputs are combined. This reduces per-device memory and compute proportionally to the TP degree — each GPU stores only its share of the KV cache (each head's K and V tensors) and processes only its heads.

Here is where the clash occurs. In MLA, all heads share a single latent representation $\mathbf{c}^{\mathrm{KV}}$. There are no per-head K and V tensors to partition — $\mathbf{c}^{\mathrm{KV}}$ is a unified compressed state. Under tensor parallelism, each GPU still needs the full $\mathbf{c}^{\mathrm{KV}}$ to compute attention for its assigned heads (after absorbing the up-projection matrices). Consequently, with TP=2, each device stores the complete 576-dimensional KV cache per token — exactly the same as with TP=1. The per-device cache is not reduced by increasing the TP degree. As the paper states:

"In contrast, Deepseek-V3 has a fixed KV cache dimension of 64 + 512 = 576, which must be fully replicated on each device regardless of the parallelism degree. This results in a higher per-device KV cache memory footprint compared to GQA-based models under the same tensor parallel configuration."

This is not a minor inefficiency. Consider a concrete comparison: LLaMA-3-70B with GQA uses 8 key-value heads, each of dimension 128, for a per-token KV cache of $2 \times 8 \times 128 = 2048$ dimensions. With TP=4, each GPU holds $2048 / 4 = 512$ dimensions per token. DeepSeek-V3 with MLA uses a fixed 576-dimensional cache that is replicated on all 4 GPUs — each device holds 576 dimensions, making MLA's per-device footprint actually larger than GQA's in this configuration. The very compression that makes MLA memory-efficient on a single GPU becomes a liability under tensor parallelism because the compression removes the structural parallelism (separate per-head KV tensors) that TP normally exploits.

Why This Problem Matters: The Decoding Memory Wall

The practical significance of this problem becomes clear when examining the inference pipeline. LLM inference splits into two phases with fundamentally different performance characteristics:

Prefilling processes the entire input prompt in one parallel forward pass. The computation is batched across all prompt tokens, so the arithmetic intensity (FLOPs per byte of memory access) is high — this phase is compute-bound.

Decoding generates output tokens one at a time, each requiring attention over the full KV cache of all previous tokens (input + generated). Each decoding step performs relatively few FLOPs — a single query vector attending over a potentially very long sequence — but must read the entire KV cache from HBM to on-chip SRAM. As context length grows, the ratio of memory reads to computations increases linearly, making this phase memory-bandwidth-bound. The decoding loop is typically the throughput bottleneck for long-context generation, and its speed is determined primarily by how fast the KV cache can be streamed from memory.

Reducing the per-device KV cache directly attacks this memory bandwidth bottleneck. If each GPU stores half the cache, it reads half the bytes per decoding step. However, as we established, MLA's unified latent representation prevents this halving from happening under standard tensor parallelism. The paper therefore identifies a structural incompatibility between MLA's compression mechanism and TP's memory distribution mechanism that prevents MLA-based models from achieving the throughput their reduced parameter counts and compressed representations otherwise promise.

Beyond throughput, this matters for deployment economics. Models like DeepSeek-V3 (685B parameters, Mixture-of-Experts) and Kimi-K2 (1T parameters) are too large to fit on a single GPU even after quantization. Tensor parallelism is not optional for these models — it is a requirement. If the KV cache cannot be distributed across TP devices, the memory consumption per device grows with context length until it exceeds GPU memory capacity, capping the usable context window. The paper's $1.79\times$ and $1.93\times$ decoding speedups at 32K context length directly translate to real serving cost reductions and higher request throughput for production deployments.

Prior Approaches and Their Limitations

The paper situates its contribution against several families of prior work, each of which addresses part of the problem but introduces its own limitations.

1. KV Cache Compression Methods (Incomplete and Lossy)

A large body of work reduces KV cache memory by discarding or altering information post-training: token pruning evicts low-importance KV entries based on attention scores or heuristics; token merging combines similar tokens into surrogate representations; cross-layer sharing reuses one cache across multiple layers; low-rank factorization approximates KV matrices via SVD at inference time; and quantization stores K/V tensors at reduced precision (int8, int4). While effective in many settings, these methods are all post-hoc approximations — they alter the model's computation from what it was trained on, inevitably introducing some degradation. As the paper notes:

"Although effective, these approaches inevitably discard or alter information in the KV cache and can degrade model performance. In contrast, TPLA leaves the KV contents intact: it reduces the amount of cache each device must hold so the model retains full information while alleviating memory pressure."

TPLA's key distinction here is that it does not compress the cache further — it distributes the already-compressed MLA cache across devices through a structural reformulation of the attention computation itself, preserving the full information content.

2. Grouped Latent Attention (GLA): The Direct Predecessor

The most directly relevant prior work is Grouped Latent Attention (GLA), proposed by Zadouri et al. (2025). GLA was designed specifically to address MLA's tensor parallelism limitation. Its approach is straightforward: partition the latent KV cache into $g$ groups (typically $g = 2$) and partition the attention heads into $g$ corresponding groups. Each group of heads attends only to its assigned latent shard. With TP=2 and $g=2$, each GPU holds half the latent dimension (e.g., $2d_h$ instead of $4d_h$) and computes attention for half the heads.

This solves the memory distribution problem. However, the paper identifies two critical limitations of GLA that motivate TPLA:

Limitation 1: Reduced representational capacity per head. In GLA, each attention head can only access $g^{-1}$ of the original latent representation. With $g=2$, each head sees only $2d_h$ dimensions rather than the full $4d_h$. The paper presents this as a fundamental expressivity loss — the model's attention heads are operating on a compressed, partial view of the KV state. This shows up dramatically in the experiments: directly converting an MLA-pretrained model to GLA (without retraining) causes WikiText-2 perplexity to explode from 6.31 to 2212 (Table 1). The information loss from forcing each head to use only half the latent dimension is catastrophic when the weights were trained expecting the full dimension.

Limitation 2: Training from scratch is required. Because GLA structurally reduces each head's access to the latent representation, an MLA-pretrained checkpoint cannot be directly converted with acceptable accuracy. Training a new GLA model from scratch would require the same massive computational investment as the original MLA model (DeepSeek-V3 reportedly used 2.788M H800 GPU hours), making GLA more a proposal for future architectures than a practical solution for existing MLA-based models. The paper frames this as a significant practical barrier:

"GLA requires training from scratch, which demands significant computational resources to validate its effectiveness."

3. Why Not Just Split Heads (Standard MLA-TP)?

A natural question is: why can't MLA simply partition heads across devices without changing the latent representation, as standard TP does for multi-head attention? The answer lies in the matrix absorption trick that makes MLA efficient during decoding.

In MLA's decoding formulation (Equations 1–2), the up-projection matrix $W^{UK}$ for keys is absorbed into the query projection, yielding an "absorbed query" $Q$. Similarly, the value up-projection $W^{UV}$ is absorbed into the output projection, yielding $W^{VO}$. The attention computation then operates directly between $Q$ and the normalized latent cache $\hat{\mathbf{c}}^{\mathrm{KV}}$, bypassing explicit key and value reconstruction. The output is computed as $O = \text{softmax}(Q(\hat{\mathbf{c}}^{\mathrm{KV}})^\top / \sqrt{d_h}) \, \hat{\mathbf{c}}^{\mathrm{KV}}$, projected through $W^{VO}$.

The critical detail: $Q$ has shape $[B, 1, h_q, 4d_h]$ — it includes all query heads. If we split heads across two GPUs (head parallelism), GPU 0 gets $Q[:,:,:h_q/2,:]$ and GPU 1 gets $Q[:,:,h_q/2:,:]$. But both GPUs still need the full $\hat{\mathbf{c}}^{\mathrm{KV}}$ of shape $[B, L, 4d_h]$ to compute their local attention scores — the head splitting doesn't reduce the KV cache that each device must load. This is precisely the problem: when the latent representation is unified (not per-head), head-level parallelism provides no memory reduction.

TPLA's insight is to split the latent dimension itself (not just the heads) across devices, so that each GPU computes attention over $2d_h$ dimensions of $\hat{\mathbf{c}}^{\mathrm{KV}}$ rather than $4d_h$, while keeping all heads active on each device through a different formulation.

How TPLA Positions Itself

TPLA occupies a specific and pragmatic point in the design space. Rather than proposing a new pretraining architecture (like GLA) or a post-hoc compression technique, it presents itself as a conversion method that transforms MLA-pretrained models into a tensor-parallel-friendly form with minimal accuracy loss and no retraining. The abstract emphasizes this drop-in compatibility:

"TPLA is drop-in compatible with models pre-trained using MLA: it supports MLA-style prefilling and enables efficient tensor-parallel decoding without retraining."

The paper's positioning rests on four claims that differentiate it from prior work:

1. Full representational capacity preserved (unlike GLA). "Unlike Grouped Latent Attention (GLA), every head in TPLA still leverages the full latent representation, maintaining stronger representational capacity." Where GLA restricts each head to a subset of the latent dimension, TPLA partitions the latent dimension across devices but arranges the computation so that each device's attention heads still operate over the complete logical latent representation — the information is distributed but not truncated.

2. No training from scratch required (unlike GLA). TPLA can directly load MLA checkpoints and convert them through matrix reparameterization (absorbing orthogonal transforms into adjacent weight matrices). The conversion process adds no learned parameters — it is a purely algebraic transformation.

3. Accuracy degradation bounded by reparameterization quality. The conversion from MLA to TPLA introduces two sources of approximation error: slicing RMSNorm across devices (each device computes normalization over $d/2$ dimensions rather than the full $d$ dimensions) and slicing the softmax attention computation (each device sees only its local attention scores, not the global sum). The paper addresses these through orthogonal transformations that rebalance information across dimensions before splitting, making the per-shard statistics approximate the global statistics. The quality of this reparameterization determines the conversion accuracy — and the paper shows that using PCA-based transforms on a small calibration set (WikiText-2) is sufficient to keep accuracy degradation modest, with an optional lightweight alignment step ($\sim 100$M tokens) providing further recovery.

4. Phase-specific optimization via prefill–decode separation. Recognizing that prefill is compute-bound while decode is memory-bound, the paper proposes using the reparameterized-but-unsliced MLA during prefill (preserving full accuracy and leveraging head-level TP for compute distribution) and switching to TPLA during decoding (with latent-dimension slicing for memory reduction). This hybrid approach, termed "TPLA (PD-sep.)," largely sidesteps the RMSNorm and softmax approximation errors for the majority of tokens (the prompt), while still achieving the memory bandwidth reduction during the bottlenecked decoding phase. Empirically, this training-free variant nearly matches original MLA accuracy on both commonsense benchmarks and LongBench (Table 1 and Table 2).

The Underlying Theoretical Tension

The paper's framing reveals a deeper insight about the relationship between compression and parallelism in attention mechanisms. MLA's compression works by merging per-head representations into a shared low-rank bottleneck — this is what achieves the KV cache reduction. But this merging simultaneously removes the natural parallelism axis (independent per-head computations) that TP relies on for memory distribution. The problem is not that MLA is bad or TP is bad — it is that MLA's compression strategy and TP's distribution strategy optimize for orthogonal goals that interact poorly.

TPLA's solution is to reintroduce a parallelism axis within the compressed representation itself — partitioning the latent dimension rather than partitioning across heads. This requires changing how attention is computed: instead of a single global softmax over all latent dimensions, each device computes a partial attention with its latent shard, and the partial outputs are summed via all-reduce. This decomposition is not mathematically equivalent to the original attention (the softmax of a sum is not the sum of softmaxes), so TPLA must absorb the approximation error through the reparameterization transforms.

The paper's key conceptual contribution is recognizing that this approximation error can be controlled through orthogonal transformations applied to the weight matrices before slicing — and that simple data-driven methods (PCA on calibration corpus activations) are sufficient to make the error small enough for practical deployment. This transforms what would be a lossy architectural change into a near-lossless conversion pipeline.

3. Technical Approach

This is primarily a systems paper with a strong theoretical reparameterization component, proposing a method to convert existing MLA-pretrained models into a form that enables efficient tensor parallelism during decoding without retraining. The core idea is to partition the latent KV representation's dimension across GPUs rather than partitioning attention heads, and to absorb orthogonal transformations into adjacent weight matrices so that per-shard RMSNorm and softmax computations approximate their global counterparts.

3.1 Reader Orientation

TPLA is a weight-reparameterization and attention-reformulation scheme that transforms a pretrained Multi-Head Latent Attention (MLA) model into a tensor-parallel-friendly form. The system takes an existing MLA checkpoint as input, applies orthogonal transformations to specific weight matrices (absorbing the transforms into adjacent projections so the model's mathematical function is unchanged), and produces a converted model where the latent KV cache dimension is split across GPUs during decoding — each GPU holds and processes only half (or $1/g$) of the latent dimensions, yet every attention head on every GPU still attends over the complete logical latent representation, with the per-shard outputs combined via all-reduce. The problem it solves is the replication bottleneck: in standard MLA under tensor parallelism, every GPU must load the full $\mathbf{c}^{\mathrm{KV}}$ latent vector, eroding MLA's memory savings. TPLA distributes this load across devices by slicing the latent dimension itself, at the cost of introducing controlled approximation errors in the per-device RMSNorm and softmax operations, which are mitigated through carefully designed orthogonal reparameterization transforms.

3.2 Big-Picture Architecture (Diagram in Words)

The TPLA system consists of four major components that transform a pretrained MLA model into an inference-optimized form:

1. Orthogonal Transform Selection Module: Given a small calibration dataset (e.g., WikiText-2), this module collects the KV latent cache activations from the pretrained MLA model, computes an orthogonal matrix $U$ (via PCA on the activation covariance, or a structured Hadamard transform) that satisfies two conditions — the squared norms of the two halves of $U$-transformed vectors must be proportional to the global norm (Condition 1, for RMSNorm slicing), and the attention score contributions from each half must approximate the global attention scores up to scaling factors (Condition 2, for softmax slicing).

2. Matrix Absorption Engine: The orthogonal transform $U$ and its transpose $U^\top$ are absorbed into the adjacent weight matrices of the MLA model: $U^\top$ is absorbed into the up-projection matrix $W^{UKV} = (W^{UK}, W^{UV})$, and $U$ is absorbed into the down-projection matrix $W^{DKV}$. The learnable RMSNorm scale parameter $W_\gamma$ is also absorbed into $W^{UKV}$. This produces new weight matrices $W^{UKV}_{\mathrm{new}} = U^\top W_\gamma W^{UKV}$ and $W^{DKV}_{\mathrm{new}} = W^{DKV} U$, while leaving the model's output mathematically equivalent under full-precision computation. After absorption, the RMSNorm operates with $\gamma = \mathbf{1}$ (identity scale), which Proposition 1 proves is necessary for the orthogonal transform to preserve RMSNorm equivalence.

3. TPLA Attention Module (Decoding): During decoding, the transformed latent vector $\mathbf{c} U$ is split along its dimension into two shards of size $2d_h$ each (for $g=2$ groups). Each GPU receives one shard and computes its local RMSNorm using only its local dimensions, with a correction factor $\sqrt{\alpha/2}$ or $\sqrt{\beta/2}$ derived from the PCA eigenvalues to approximate the global RMS value. Each GPU then computes local attention scores between its shard of $Q$ and its shard of $\hat{\mathbf{c}}^{\mathrm{KV}}$, adds the (replicated) RoPE positional scores, applies softmax locally, computes the attention output, and projects through its shard of $W^{VO}$. The two GPUs' outputs are summed via an all-reduce operation. The full query $Q$ is replicated across devices (all heads present on each GPU), while the latent dimension is partitioned — this is the opposite of standard head-wise TP.

4. Prefill-Decode Separation Controller: During the compute-bound prefilling phase, the system uses the reparameterized MLA form (with absorbed orthogonal transforms) but does not slice RMSNorm or softmax — the computation proceeds identically to the original MLA, with head-level tensor parallelism for compute distribution. The resulting KV cache is stored in the reparameterized latent space. During the memory-bound decoding phase, the system switches to TPLA mode with sliced RMSNorm and softmax, consuming the prefill KV cache directly (since the reparameterization ensures the latent representations are compatible). This avoids approximation error for the majority of tokens (the prompt) while still achieving memory bandwidth reduction during the throughput-critical decoding loop.

Information flows as follows: pretrained MLA checkpoint → orthogonal transform selection (PCA on calibration activations or Hadamard construction) → matrix absorption ($U$ into $W^{DKV}$, $U^\top$ and $W_\gamma$ into $W^{UKV}$) → reparameterized model → at inference: prefill with unsliced MLA (head TP) → KV cache stored in reparameterized latent space → decode switch to TPLA (latent dimension sliced across GPUs, local RMSNorm/softmax with correction factors, all-reduce outputs).

3.3 Roadmap for the Deep Dive

• First, the matrix absorption formulation in MLA — what gets absorbed where, and why this absorption is the foundation that makes TPLA's weight reparameterization possible. Without understanding absorption, the orthogonal transform insertion appears arbitrary.
• Second, the core TPLA computation and its structural differences from GLA — showing exactly how TPLA partitions the latent dimension while keeping all heads active, and why this formulation is algebraically equivalent to a GLA system with doubled heads (enabling FlashAttention-3 compatibility).
• Third, the RMSNorm slicing problem and its solution — Proposition 1 (equivalence under orthogonal transforms), Condition 1 (the proportionality requirement for per-shard norm approximation), and the derivation of correction factors $\alpha$ and $\beta$ from PCA eigenvalues.
• Fourth, the softmax slicing problem and its solution — Condition 2 (the proportionality requirement for per-shard attention scores), why softmax is harder than RMSNorm (the exponential amplifies asymmetries), and how PCA addresses this while Hadamard does not.
• Fifth, the two reparameterization methods (Hadamard and PCA) — their mathematical constructions, how they satisfy Conditions 1 and 2, and their contrasting failure modes (Hadamard balances norms but fails on softmax; PCA concentrates information but loses effectiveness with $g > 2$).
• Sixth, the prefill–decode separation strategy — why prefill is compute-bound and benefits from head-level TP without slicing, while decode is memory-bound and benefits from latent-dimension slicing, and how the reparameterization enables seamless KV cache reuse across phases.
• Seventh, the complexity analysis — a detailed comparison of FLOPs, memory traffic, and communication patterns between MLA-TP and TPLA-TP, showing where TPLA's computational overhead comes from (query replication) and why it is dominated by memory bandwidth savings at long context lengths.

3.4 Detailed, Sentence-Based Technical Breakdown

MLA Matrix Absorption: The Foundation for Reparameterization

MLA achieves its decoding efficiency through a series of matrix absorption steps that eliminate the need to explicitly reconstruct full-size key and value tensors from the compressed latent representation. Understanding these absorptions is essential because TPLA's reparameterization strategy inserts orthogonal transforms into the absorbed weight matrices, not into the original unabsorbed ones.

The unabsorbed MLA computation (Equation 4 in the paper) proceeds as follows. The input hidden states $X \in \mathbb{R}^{B \times L \times D}$ are projected through learned down-projection matrices to produce two compressed representations: the KV latent $\mathbf{c}^{\mathrm{KV}} = X W^{DKV}$ where $W^{DKV} \in \mathbb{R}^{D \times 4d_h}$, and the query latent $\mathbf{c}^{Q} = X W^{DQ}$ where $W^{DQ} \in \mathbb{R}^{D \times r_q}$. The KV latent is normalized via RMSNorm to produce $\hat{\mathbf{c}}^{\mathrm{KV}} = \mathrm{RMSNorm}(\mathbf{c}^{\mathrm{KV}})$. Full-size multi-head keys and values are then reconstructed by up-projection: $\mathbf{k} = \hat{\mathbf{c}}^{\mathrm{KV}} W^{UK}$ and $\mathbf{v} = \hat{\mathbf{c}}^{\mathrm{KV}} W^{UV}$, where $W^{UK}, W^{UV} \in \mathbb{R}^{4d_h \times (h_q \cdot d_h)}$ map the $4d_h$-dimensional compressed representation to the $h_q \cdot d_h$-dimensional concatenated per-head key/value space. Similarly, full-size queries are reconstructed as $\mathbf{q} = \mathbf{c}^{Q} W^{UQ}$, where $W^{UQ} \in \mathbb{R}^{r_q \times (h_q \cdot d_h)}$. The attention output is $\tilde{O} = \mathrm{softmax}(\mathbf{q} \mathbf{k}^\top / \sqrt{d_h + d_r} + \text{positional}) \cdot \mathbf{v} \cdot W^{O}$.

The key absorption (Equation 5). During decoding (where the query sequence length $L_q = 1$), the matrix $W^{UK}$ can be factored into the query computation rather than being applied to the KV cache:

$O W^{O} = \mathrm{softmax}(\mathbf{q} (\hat{\mathbf{c}}^{\mathrm{KV}} W^{UK})^\top) \hat{\mathbf{c}}^{\mathrm{KV}} W^{UV} W^{O}$

where $O$ is the pre-output attention tensor, $W^{O} \in \mathbb{R}^{(h_q \cdot d_h) \times D}$ is the output projection, and the other symbols are as defined above.

What it computes: The key up-projection $W^{UK}$ is moved from the key side (where it would need to be applied to the growing KV cache) to the query side (where it is applied once to the single new query token). This is possible because matrix multiplication is associative: $\mathbf{q} (W^{UK^\top} \hat{\mathbf{c}}^{\mathrm{KV}^\top}) = (\mathbf{q} W^{UK^\top}) \hat{\mathbf{c}}^{\mathrm{KV}^\top}$. The term $\mathbf{q} W^{UK^\top}$ can be precomputed once per decoding step, producing an "absorbed query" $Q$ that directly attends to the raw normalized latent $\hat{\mathbf{c}}^{\mathrm{KV}}$.

Why this form: Without absorption, each decoding step would require computing $\hat{\mathbf{c}}^{\mathrm{KV}} W^{UK}$ for the entire KV cache — an $O(L \cdot 4d_h \cdot h_q d_h)$ operation that grows with context length. Absorption reduces this to a one-time $O(1 \cdot r_q \cdot h_q d_h)$ query precomputation plus $O(L \cdot 4d_h)$ attention score computation, eliminating the per-token up-projection cost and the need to cache full-size keys.

The value-output absorption (Equation 5 continued). The same associativity principle allows absorbing $W^{UV}$ into the output projection:

$\mathrm{softmax}(Q (\hat{\mathbf{c}}^{\mathrm{KV}})^\top) \hat{\mathbf{c}}^{\mathrm{KV}} W^{VO}$

where $Q = \mathbf{q} (W^{UK})^\top \in \mathbb{R}^{B \times 1 \times h_q \times 4d_h}$ is the absorbed query (with $W^{UK}$ absorbed into the query projection), and $W^{VO} = W^{UV} W^{O} \in \mathbb{R}^{(h_q \cdot 4d_h) \times D}$ is the combined value-output projection matrix.

What it computes: The attention-weighted sum over the latent cache $\hat{\mathbf{c}}^{\mathrm{KV}}$ (weighted by softmax scores) is projected directly to the output dimension $D$ through $W^{VO}$, bypassing explicit value reconstruction and separate output projection. The attention operates in the compressed $4d_h$-dimensional space rather than the full $h_q \cdot d_h$-dimensional head space.

Why this form: The paper notes that in practice $W^{UV}$ is typically not absorbed into $W^{O}$ to avoid generating an impractically large matrix — $W^{VO}$ has shape $(h_q \cdot 4d_h) \times D$, which for DeepSeek-V3 with $h_q = 128$ and $d_h = 128$ would be $65536 \times D$, much larger than the separate $W^{UV}$ (shape $512 \times 16384$) and $W^{O}$ (shape $16384 \times D$). The absorption is a conceptual tool for understanding the decoding formulation, but the actual implementation keeps these matrices separate where the dimensions would otherwise explode.

The RoPE handling. Rotary Position Embedding (RoPE) is applied only to a small decoupled component: a separate query $\mathbf{q}^{\mathrm{PE}}$ (dimension $d_r$ per head, typically 64) and a separate key $\mathbf{k}^{\mathrm{PE}}$ (also dimension $d_r$, shared across heads). These carry the position information while the main $4d_h$-dimensional latent attention handles content. The final attention score is the sum:

$\mathrm{softmax}\left(\frac{Q (\hat{\mathbf{c}}^{\mathrm{KV}})^\top + \mathbf{q}^{\mathrm{PE}} (\mathbf{k}^{\mathrm{PE}})^\top}{\sqrt{d_h + d_r}}\right)$

This decoupling is critical for TPLA because the RoPE positional component is not latent-compressed — $\mathbf{k}^{\mathrm{PE}}$ is generated by a separate projection $W^{KR} \in \mathbb{R}^{D \times d_r}$ and must be replicated across all devices in TPLA, since each device needs the full positional attention scores to add to its local content-based scores. The paper explicitly notes this replication: "the shard of the key positional embedding $\mathbf{k}^{\mathrm{PE}}$ must be replicated across devices so that the local positional values remain consistent with the global values" (Section 4.2).

The Core TPLA Computation and Structural Differences from GLA

TPLA's attention computation (Equations 6–7) represents a fundamental redesign of how tensor parallelism interacts with the latent representation. Where standard MLA-TP splits attention heads across devices (each GPU gets $h_q/2$ heads) but keeps the latent dimension intact (each GPU loads the full $4d_h$-dimensional cache), TPLA does the opposite: it keeps all heads on each device (each GPU has the full $h_q$ heads) but splits the latent dimension (each GPU loads a $2d_h$-dimensional shard of the cache, for $g=2$).

The computation on GPU 0 (for $g=2$, TP degree $k=2$) is:

$O_{\mathbf{0}} = \mathrm{softmax}\!\left(\frac{Q_{\mathbf{0}} (\hat{\mathbf{c}}_{\mathbf{0}}^{\mathrm{KV}})^\top}{\sqrt{d_h}}\right) \hat{\mathbf{c}}_{\mathbf{0}}^{\mathrm{KV}} \in \mathbb{R}^{B \times 1 \times h_q \times 2d_h}$

$\tilde{O_{\mathbf{0}}} = O_{\mathbf{0}} W_{\mathbf{0}}^{VO} \in \mathbb{R}^{B \times 1 \times D}$

where $Q_{\mathbf{0}} \in \mathbb{R}^{B \times 1 \times h_q \times 2d_h}$ is the first half of the absorbed query tensor (split along the feature dimension, not the head dimension), $\hat{\mathbf{c}}_{\mathbf{0}}^{\mathrm{KV}} \in \mathbb{R}^{B \times L \times 2d_h}$ is the first half of the normalized latent cache, $W_{\mathbf{0}}^{VO} \in \mathbb{R}^{(h_q \cdot 2d_h) \times D}$ is the first half of the combined value-output projection (split along the input dimension), and $d_h$ is the per-head dimension used in the softmax scaling (using the original $d_h$ rather than $2d_h$ preserves the temperature of the original attention).

What it computes: GPU 0 computes attention scores between all $h_q$ query heads and the first $2d_h$ dimensions of the latent cache, applies softmax to these partial scores, computes a partial attention-weighted sum over those $2d_h$ cache dimensions, and projects the result to the output space. GPU 1 does the same with the second $2d_h$ dimensions. The two partial outputs are summed: $O = \mathrm{AllReduce}(\tilde{O_{\mathbf{0}}} + \tilde{O_{\mathbf{1}}})$.

Why this form: The key design choice is that every attention head sees the full $4d_h$-dimensional latent representation — just split across two GPUs rather than truncated. Head $i$ on GPU 0 attends to dimensions 0 through $2d_h-1$ of the latent cache, while head $i$ on GPU 1 attends to dimensions $2d_h$ through $4d_h-1$. The complete attention score for head $i$ across the full latent dimension is the sum of the two partial scores (up to the softmax decomposition error). This preserves the representational capacity that GLA loses — in GLA, head $i$ on GPU 0 would only ever see $2d_h$ dimensions, period, with no pathway to recover the full $4d_h$ information. In TPLA, the information is distributed but not discarded; the all-reduce sum reconstructs (an approximation of) the full attention output.

The softmax non-linearity problem. The decomposition is not exact because:

$\mathrm{softmax}(a + b) \neq \mathrm{softmax}(a) + \mathrm{softmax}(b)$

where $a = Q_{\mathbf{0}}(\hat{\mathbf{c}}_{\mathbf{0}}^{\mathrm{KV}})^\top$ and $b = Q_{\mathbf{1}}(\hat{\mathbf{c}}_{\mathbf{1}}^{\mathrm{KV}})^\top$ are the per-shard score contributions for a given head and token. The correct global attention output would be:

$\mathrm{softmax}(a + b + \text{pos}) \cdot [\hat{\mathbf{c}}_{\mathbf{0}}^{\mathrm{KV}}; \hat{\mathbf{c}}_{\mathbf{1}}^{\mathrm{KV}}]$

but TPLA computes:

$\mathrm{softmax}(a + \text{pos}) \cdot \hat{\mathbf{c}}_{\mathbf{0}}^{\mathrm{KV}} + \mathrm{softmax}(b + \text{pos}) \cdot \hat{\mathbf{c}}_{\mathbf{1}}^{\mathrm{KV}}$

The approximation error depends on how much the softmax normalization is distorted by computing it separately on each shard rather than jointly. This is why Conditions 1 and 2 are needed — they ensure that the $U$-transformed activations have statistical properties that make the decomposition error small.

TPLA as a special case of GLA (Section 4.4). The paper establishes an important equivalence that enables TPLA to leverage existing optimized attention kernels. If we conceptually duplicate the query tensor along the head dimension — creating a $Q'$ with $2h_q$ heads where the first $h_q$ heads are identical to the second $h_q$ heads — then TPLA's computation becomes algebraically equivalent to GLA with $2h_q$ heads and a full $4d_h$ latent dimension, where the query tensor is partitioned as:

$Q' = \begin{pmatrix} Q_{\mathbf{0,0}} & Q_{\mathbf{0,1}} \\ Q_{\mathbf{1,0}} & Q_{\mathbf{1,1}} \\ Q_{\mathbf{0,0}} & Q_{\mathbf{0,1}} \\ Q_{\mathbf{1,0}} & Q_{\mathbf{1,1}} \end{pmatrix}$

and after GLA-style sharding (splitting both heads and features), GPU 0 gets the concatenation $[Q_{\mathbf{0,0}} \; Q_{\mathbf{0,1}}] \in \mathbb{R}^{B \times L \times h_q \times 4d_h}$ and GPU 1 gets $[Q_{\mathbf{1,0}} \; Q_{\mathbf{1,1}}]$. The crucial difference from actual GLA (which would only get the diagonal blocks $Q_{\mathbf{0,0}}$ and $Q_{\mathbf{1,1}}$) is that TPLA's duplication gives each GPU access to all $4d_h$ latent dimensions through its full set of heads, while GLA restricts each GPU to $2d_h$ latent dimensions through half the heads.

What this means practically: TPLA can be implemented using FlashAttention-3 with a GLA-compatible sharding pattern, simply by providing the duplicated query tensor. The computational overhead is the cost of computing attention for twice as many effective heads, but the memory savings (halving the per-device KV cache) dominate at long context lengths. The paper analyzes this tradeoff in Section 4.5: for TP=2, the main attention complexity is $\mathcal{O}(L_q \times S_{KV} \times h_q \times 2d_h \times 2)$ for TPLA versus $\mathcal{O}(L_q \times S_{KV} \times h_q/2 \times 4d_h \times 2)$ for MLA-TP — arithmetically equivalent for the core attention, but with additional overhead from the duplicated RoPE computation and the replicated query projection.

RMSNorm Slicing: The Problem and Solution

RMSNorm (Root Mean Square Normalization) is the normalization layer applied to the KV latent $\mathbf{c}^{\mathrm{KV}}$ before attention. Its standard form (Equations 8–9) for an input vector $\mathbf{x} \in \mathbb{R}^d$ is:

$\text{RMS}(\mathbf{x}) = \sqrt{\frac{1}{d} \sum_{i=1}^{d} x_i^2 + \epsilon} = \sqrt{\frac{1}{d} \|\mathbf{x}\|_2^2 + \epsilon}$

$\text{RMSNorm}(\gamma, \mathbf{x}) = \frac{\mathbf{x}}{\text{RMS}(\mathbf{x})} \odot \gamma$

where $\epsilon$ is a small constant for numerical stability, $\gamma \in \mathbb{R}^d$ is a learned per-dimension scale parameter, and $\|\mathbf{x}\|_2^2 = \sum_{i=1}^d x_i^2$ denotes the squared L2 norm.

What it computes: For each token in the sequence, the RMSNorm divides every element of the input vector by the root-mean-square of all elements in that vector, then multiplies by a learned scale per dimension. This normalizes the vector to have unit RMS (before scaling), stabilizing training and inference by controlling activation magnitudes.

Why this form: RMSNorm uses the root-mean-square rather than the standard deviation (as in LayerNorm) because it avoids computing and subtracting the mean — this is computationally cheaper and has been shown to work equally well in practice for transformer architectures. The learnable scale $\gamma$ allows the model to recover any desired per-dimension variance after normalization.

The slicing problem. When TPLA splits $\mathbf{c}^{\mathrm{KV}} \in \mathbb{R}^{4d_h}$ into two shards $\mathbf{c}_{\mathbf{0}}^{\mathrm{KV}} \in \mathbb{R}^{2d_h}$ and $\mathbf{c}_{\mathbf{1}}^{\mathrm{KV}} \in \mathbb{R}^{2d_h}$ on different GPUs, each GPU can only compute the RMS using its $2d_h$ local dimensions:

$\text{RMS}_{\text{local}}(\mathbf{c}_{\mathbf{0}}) = \sqrt{\frac{1}{2d_h} \|\mathbf{c}_{\mathbf{0}}\|_2^2 + \epsilon}$

But the true RMSNorm requires the global RMS over all $4d_h$ dimensions:

$\text{RMS}_{\text{global}}(\mathbf{c}) = \sqrt{\frac{1}{4d_h} (\|\mathbf{c}_{\mathbf{0}}\|_2^2 + \|\mathbf{c}_{\mathbf{1}}\|_2^2) + \epsilon}$

These are only equal if $\|\mathbf{c}_{\mathbf{0}}\|_2^2 = \|\mathbf{c}_{\mathbf{1}}\|_2^2$ — that is, if the norm is equally distributed across the two halves. For an arbitrary pretrained MLA model, there is no reason this equality holds; the latent dimensions may have very different typical magnitudes.

Proposition 1: Orthogonal transform equivalence. The paper proves that applying any orthogonal matrix $U \in \mathbb{R}^{d \times d}$ (satisfying $U U^\top = \mathbf{I}$) to the input before RMSNorm has an equivalent reparameterized form:

$\text{RMSNorm}(\mathbf{1}, \mathbf{c}) = \text{RMSNorm}(\mathbf{1}, \mathbf{c} U) U^\top$

where $\text{RMSNorm}(\mathbf{1}, \cdot)$ denotes RMSNorm with identity scale ($\gamma = \mathbf{1}$). The proof relies on three facts: (1) orthogonal transformations preserve the L2 norm ($\|\mathbf{c} U\|_2^2 = \|\mathbf{c}\|_2^2$), so the RMS value is unchanged; (2) RMSNorm can be expressed as matrix multiplication $\text{RMSNorm}(\gamma, \mathbf{c}) = D_c \mathbf{c} W_\gamma$ where $D_c = \text{diag}(1/\text{RMS}(c_1), \ldots, 1/\text{RMS}(c_L))$ is a diagonal matrix of per-token inverse RMS values and $W_\gamma = \text{diag}(\gamma_1, \ldots, \gamma_d)$ is a diagonal matrix of per-dimension scales; (3) when $W_\gamma = \mathbf{I}$, the $U$ and $U^\top$ cancel because $D_{cU} = D_c$ (norm preserved) and $U \mathbf{I} U^\top = \mathbf{I}$.

What this enables: The orthogonal transform can be "pushed through" the RMSNorm — applying $U$ before RMSNorm and $U^\top$ after produces the identical output. This means we can pre-transform the latent activations with $U$ without changing the model's output, as long as we also apply $U^\top$ to the subsequent computation (which gets absorbed into $W^{UKV}$).

Absorbing the scale parameter $\gamma$. The paper notes that Proposition 1 only holds when $\gamma = \mathbf{1}$. To satisfy this, the learned scale $W_\gamma$ is absorbed into the up-projection matrix:

$W^{UKV}_{\mathrm{new}} = U^\top W_\gamma W^{UKV}$

This is a critical step: without absorbing $W_\gamma$, the orthogonal equivalence would fail. The absorption modifies the up-projection weights to incorporate both the inverse orthogonal transform $U^\top$ and the per-dimension RMSNorm scale $\gamma$ simultaneously. The down-projection $W^{DKV}$ absorbs the forward transform: $W^{DKV}_{\mathrm{new}} = W^{DKV} U$.

Condition 1: The proportionality requirement (Equation 18). After the orthogonal transform, we need the norm of each half to be proportional to the global norm, so that local RMS computations can estimate the global RMS:

$\alpha \|(\mathbf{c} U)_{\mathbf{0}}\|^2_2 \approx \beta \|(\mathbf{c} U)_{\mathbf{1}}\|^2_2 \approx \|\mathbf{c} U\|^2_2 = \|\mathbf{c}\|^2_2$

where $\alpha$ and $\beta$ are fixed constants (not depending on the specific input $\mathbf{c}$), $(\mathbf{c} U)_{\mathbf{0}}$ is the first half of the transformed vector, and $(\mathbf{c} U)_{\mathbf{1}}$ is the second half.

What it computes: The sum of the per-shard norms (scaled by their respective proportionality constants $\alpha$ and $\beta$) should equal the global norm. In practice, $\alpha$ and $\beta$ are chosen so that $\alpha \|(\mathbf{c} U)_{\mathbf{0}}\|^2_2$ approximates the global squared norm. Then each GPU can estimate the global RMS from its local shard:

$\text{RMS}(\mathbf{c}) \approx \sqrt{\frac{\alpha}{2}} \cdot \text{RMS}_{\text{local}}((\mathbf{c} U)_{\mathbf{0}})$

where the $\sqrt{\alpha/2}$ factor accounts for the ratio of full dimension ($d$) to half-dimension ($d/2$) times the proportionality constant.

Derivation of $\alpha$ and $\beta$ from PCA (Equation 25). When using PCA as the orthogonal transform $U$, the eigenvectors are ordered by decreasing eigenvalue. The eigenvalue $\lambda_i$ equals the variance of the activations along the $i$-th principal component (assuming mean-centered data). The proportion of total variance captured by the first $d/2$ components is:

$\alpha = \frac{\sum_{i=1}^{d/2} \lambda_i}{\sum_{i=1}^{d} \lambda_i}, \quad \beta = \frac{\sum_{i=d/2}^{d} \lambda_i}{\sum_{i=1}^{d} \lambda_i}$

where $d = 4d_h$ is the total latent dimension, $\lambda_i$ are the eigenvalues of the activation covariance matrix $\Sigma_F = U \Lambda U^\top$ computed from the calibration dataset activations $F \in \mathbb{R}^{(B \cdot L) \times d}$.

What this means physically: PCA rotates the activation space so that the first $d/2$ dimensions capture exactly fraction $\alpha$ of the total variance, and the remaining $d/2$ dimensions capture fraction $\beta = 1 - \alpha$. For typical neural network activations, PCA eigenvalues decay rapidly — the first few components capture most of the variance — so $\alpha$ is close to 1 and $\beta$ is close to 0. The first shard contains almost all the "energy," and its local RMS (scaled by $\sqrt{\alpha/2}$) closely approximates the global RMS. The second shard has very little energy, so its contribution to the normalization is small (and the scaling factor $\sqrt{\beta/2}$ correctly accounts for this).

Why Hadamard works for RMSNorm. The Hadamard transform spreads information uniformly: after applying $H_d / \sqrt{d}$ to any vector, the elements tend to have similar magnitudes. This means $\|(\mathbf{c} H)_{\mathbf{0}}\|^2_2 \approx \|(\mathbf{c} H)_{\mathbf{1}}\|^2_2 \approx \frac{1}{2} \|\mathbf{c}\|^2_2$, so $\alpha = \beta = 2$ (since $\alpha \cdot \frac{1}{2}\|\mathbf{c}\|^2 = \|\mathbf{c}\|^2$). Each shard captures approximately half the norm, and the correction factor $\sqrt{\alpha/2} = \sqrt{2/2} = 1$ — no correction needed, since each half has the same per-dimension RMS as the full vector. The paper confirms this empirically in Figure 2, where "TPLA (norm only)" with Hadamard achieves performance comparable to the original MLA.

Softmax Slicing: The Harder Problem

Softmax slicing is more challenging than RMSNorm slicing because of the non-linearity and the exponential amplification of differences. The global attention score for a single query token attending to a single key token, for a given head, is:

$s = \frac{Q (\hat{\mathbf{c}}^{\mathrm{KV}})^\top}{\sqrt{d_h}} + \frac{\mathbf{q}^{\mathrm{PE}} (\mathbf{k}^{\mathrm{PE}})^\top}{\sqrt{d_h + d_r}}$

where the first term is the content-based score (summing over $4d_h$ latent dimensions) and the second term is the positional score.

In TPLA, each GPU computes only its local content score $s_{\mathbf{0}} = Q_{\mathbf{0}} (\hat{\mathbf{c}}_{\mathbf{0}}^{\mathrm{KV}})^\top / \sqrt{d_h}$ (on GPU 0) or $s_{\mathbf{1}} = Q_{\mathbf{1}} (\hat{\mathbf{c}}_{\mathbf{1}}^{\mathrm{KV}})^\top / \sqrt{d_h}$ (on GPU 1). The global content score is $s = s_{\mathbf{0}} + s_{\mathbf{1}}$ (plus positional). The softmax on GPU 0 is computed as:

$p_{\mathbf{0}, i} = \frac{\exp(s_{\mathbf{0}, i} + \text{pos}_i)}{\sum_j \exp(s_{\mathbf{0}, j} + \text{pos}_j)}$

whereas the true global softmax would be:

$p_i = \frac{\exp(s_{\mathbf{0}, i} + s_{\mathbf{1}, i} + \text{pos}_i)}{\sum_j \exp(s_{\mathbf{0}, j} + s_{\mathbf{1}, j} + \text{pos}_j)}$

These are not equal unless $s_{\mathbf{1}, i} = s_{\mathbf{1}, j}$ for all $i, j$ (i.e., the missing shard contributes a constant offset, which is absorbed into the softmax normalization) or unless one shard dominates ($s_{\mathbf{0}} \gg s_{\mathbf{1}}$ everywhere, so $\exp(s_{\mathbf{0}} + s_{\mathbf{1}}) \approx \exp(s_{\mathbf{0}}) \cdot \text{const}$).

Condition 2: The proportionality requirement (Equation 22). After orthogonal transformation, the global content-based score should be approximable from either local shard up to a multiplicative constant:

$Q' (\hat{\mathbf{c}}^{\mathrm{KV}} U)^\top \approx \mu \cdot Q'_{\mathbf{0}} (\hat{\mathbf{c}}^{\mathrm{KV}} U)_{\mathbf{0}}^\top \approx \nu \cdot Q'_{\mathbf{1}} (\hat{\mathbf{c}}^{\mathrm{KV}} U)_{\mathbf{1}}^\top$

where $Q' = Q U$ is the transformed query (with $U$ absorbed from Equation 21), $Q'_{\mathbf{0}}$ and $Q'_{\mathbf{1}}$ are its two halves along the feature dimension, and $\mu, \nu$ are scaling constants.

What it requires: The contribution of one latent shard to the attention score must be approximately proportional to the global attention score across all query-key pairs. This is a stronger condition than Condition 1 — it requires not just that norms are balanced, but that the dot products between queries and keys, when restricted to one shard, are proportional to the full dot products. In the extreme, if $Q'_{\mathbf{0}}$ and $(\hat{\mathbf{c}}^{\mathrm{KV}} U)_{\mathbf{0}}$ contain only noise uncorrelated with the full dot product, the local softmax would be essentially random.

Why Hadamard fails for softmax. The paper provides a concrete counterexample (Section 4.3.1). Consider $Q = (100, 0, 0, 0)$ and $\mathbf{c} = (0, 0, 80, 0)$. The true dot product is $0$ (they are non-zero in different dimensions). After Hadamard transform with $H_4$:

$Q' = (50, 50, 50, 50), \quad \mathbf{c}' = (40, 40, -40, -40)$

The element-wise product $Q' \odot \mathbf{c}' = (2000, 2000, -2000, -2000)$, so the full dot product is 0 (correct). But split into halves: first half sum = $2000 + 2000 = 4000$, second half sum = $-2000 + (-2000) = -4000$. Neither half equals 0, and they don't satisfy proportionality ($4000 \neq \mu \cdot 0$ for any finite $\mu$). The Hadamard transform creates both positive and negative contributions that cancel globally but not locally, causing large per-shard softmax errors. The exponential in softmax amplifies these errors — token $i$ might get a high score on GPU 0 and a low score on GPU 1, and the separate softmaxes produce inconsistent probability distributions.

Why PCA works for softmax. PCA orders dimensions by variance. The first $d/2$ principal components capture fraction $\alpha$ of the total variance in both queries and keys (since the same $U$ is applied to both, via absorption). This means the dot product $Q' (\hat{\mathbf{c}}^{\mathrm{KV}} U)^\top$ is dominated by the first $d/2$ dimensions — the contribution from the second $d/2$ dimensions is small relative to the first. Therefore:

$Q' (\hat{\mathbf{c}}^{\mathrm{KV}} U)^\top \approx Q'_{\mathbf{0}} (\hat{\mathbf{c}}^{\mathrm{KV}} U)_{\mathbf{0}}^\top$

with $\mu \approx 1 / \alpha$ (since the first half captures fraction $\alpha$ of the dot product variance). The second half contributes a small perturbation that adds noise but does not fundamentally distort the softmax distribution. The paper sets $\mu = \alpha$ and $\nu = \beta$ (Equation 25, "the metrics $\mu$ and $\nu$ are defined in the same manner, making them equivalent to $\alpha$ and $\beta$, respectively"), meaning each GPU scales its local attention scores by the fraction of variance it captures before computing softmax. This scaling adjusts the softmax temperature per device so that the relative magnitudes of scores are approximately correct.

The empirical evidence for PCA's softmax performance. Figure 2 shows that "TPLA (softmax only)" with PCA-based reparameterization substantially outperforms the Hadamard variant. The paper hypothesizes: "the exponential nature of softmax makes it more sensitive to imbalance. Although Hadamard-based reparameterization achieves statistical balance across devices, small per-sample perturbations may result in significant asymmetries." This is a crucial insight: statistical balance in expectation ($\mathbb{E}[\|(\mathbf{c} U)_{\mathbf{0}}\|^2] = \mathbb{E}[\|(\mathbf{c} U)_{\mathbf{1}}\|^2]$) is not sufficient for softmax — what matters is per-sample balance in the attention scores, and the Hadamard transform's mixed-sign structure creates large per-sample variance even though the mean is balanced.

Reparameterization Methods: Hadamard vs. PCA

The paper explores two orthogonal transform constructions, each with different properties regarding the two slicing conditions.

Hadamard Matrix Transformation (Section 4.3.1). A Hadamard matrix $H_d \in \mathbb{R}^{d \times d}$ (of Sylvester/Walsh type) is constructed recursively:

$H_{2n} = \begin{pmatrix} H_n & H_n \\ H_n & -H_n \end{pmatrix}, \quad H_1 = (1)$

where each entry is either $+1$ or $-1$. The matrix satisfies $H_d H_d^\top = d \cdot \mathbf{I}$, so the normalized version $H_d / \sqrt{d}$ is orthogonal. A random diagonal sign matrix $D$ (entries $\pm 1$) is multiplied with $H_d$ to break deterministic structure while preserving orthogonality.

What it does: The Hadamard transform mixes all input dimensions uniformly — each output element is a sum of all input elements with equal-magnitude coefficients (signs arranged according to the Hadamard pattern). This "spreads out" concentrated activations: a sparse vector like $(100, 0, 0, 0)$ becomes $(50, 50, 50, 50)$ after $H_4/2$, distributing energy equally across all dimensions.

Why it satisfies Condition 1 (RMSNorm): Because the energy is balanced, $\|(\mathbf{c} H)_{\mathbf{0}}\|^2_2 \approx \|(\mathbf{c} H)_{\mathbf{1}}\|^2_2$ for any input, so $\alpha = \beta = 2$ and no correction is needed beyond the dimension scaling. The paper shows this works extremely well empirically.

Why it fails Condition 2 (softmax): The mixed-sign structure means that while magnitudes are balanced, the signs of contributions are not — the dot product $Q_{\mathbf{0}} (\hat{\mathbf{c}}_{\mathbf{0}})^\top$ can be positive when the global dot product is zero (due to cancellation with the negative contributions in the second half), as shown in the counterexample. This per-sample sign asymmetry, amplified by the exponential, destroys softmax accuracy.

PCA-Based Transformation (Section 4.3.2). The PCA approach uses the eigenvectors of the activation covariance matrix as the orthogonal transform $U$. The procedure is:

1. Collect KV latent cache activations $F \in \mathbb{R}^{(B \cdot L) \times d}$ from the pretrained MLA model on a calibration dataset (WikiText-2 is used in experiments).
2. Center the data (mean subtraction implied by the covariance computation).
3. Compute the covariance matrix $\Sigma_F = \frac{1}{n} F^\top F$ and its eigenvalue decomposition $\Sigma_F = U \Lambda U^\top$, where $U \in \mathbb{R}^{d \times d}$ is the matrix of eigenvectors (orthogonal) and $\Lambda = \text{diag}(\lambda_1, \ldots, \lambda_d)$ contains eigenvalues in descending order.
4. Use $U$ as the orthogonal transform: $\mathbf{c} \rightarrow \mathbf{c} U$ projects activations onto the principal components, with the first dimension capturing the most variance and the last dimension capturing the least.

What it does: PCA rotates the activation space so that dimensions are ordered by importance (variance). The first $d/2$ dimensions capture most of the "signal," while the last $d/2$ dimensions capture mostly noise and minor variations. When the latent dimension is split, GPU 0 gets the high-variance dimensions and GPU 1 gets the low-variance dimensions.

Why it satisfies Condition 1: The variance captured by each half is exactly the sum of its eigenvalues, so $\alpha = \sum_{i=1}^{d/2} \lambda_i / \sum_{i=1}^d \lambda_i$ is the natural proportionality constant. GPU 0's local RMS, scaled by $\sqrt{\alpha/2}$, closely approximates the global RMS because the low-variance shard on GPU 1 contributes negligibly to the norm.

Why it satisfies Condition 2: The dot product between query and key activations, when both are projected onto the PCA basis, is dominated by the high-variance components. The contribution from the first $d/2$ dimensions is approximately $\alpha$ times the full dot product (in expectation over the data distribution). GPU 0's local attention scores, scaled by $1/\alpha$ (or equivalently, computing softmax with temperature adjusted), approximate the global attention scores because the missing second shard is small and approximately uncorrelated noise. GPU 1's contribution can be treated similarly with scaling factor $\nu = \beta$, though since $\beta \ll 1$ typically, it is often negligible.

The calibration dataset choice. The paper uses WikiText-2 and SmolLM-Corpus for PCA calibration. The quality of the PCA transform depends on how representative the calibration activations are of the deployment data distribution — if the model encounters very different activation patterns at test time, the PCA-derived $\alpha$ and $\beta$ may not accurately reflect the variance partitioning, leading to larger approximation errors. The paper does not extensively study this distribution shift effect, though the strong LongBench results suggest the PCA transform generalizes reasonably across domains.

The $g > 2$ limitation. The paper explicitly notes that PCA "concentrates most of the data's informative content in the first few dimensions" and that "when $g > 2$, it probably fails to maintain effectiveness." If the latent dimension is partitioned into more than two groups, the middle groups would contain moderate-variance dimensions that are neither dominant enough to approximate the global computation nor negligible enough to ignore. The Hadamard transform, by contrast, would continue to balance all groups equally, making it potentially better for higher TP degrees — but only if the softmax slicing problem can be solved. The paper flags "optimized Hadamard-like orthogonal matrices to balance softmax slicing" as future work.

Prefill–Decode Separation: Phase-Specific Optimization

The separation strategy (Section 4.5) exploits the fundamentally different computational characteristics of the two inference phases to minimize accuracy loss while maximizing throughput.

Why prefill is compute-bound. During prefilling, the entire input prompt (potentially thousands of tokens) is processed in one batch. The attention computation involves $O(L_q^2)$ operations where $L_q$ is the prompt length — quadratic in sequence length. The arithmetic intensity (FLOPs per byte of memory access) is high because the $O(L_q^2)$ attention operations amortize the $O(L_q)$ memory reads of the model weights and the growing KV cache. The bottleneck is the GPU's computational throughput (FLOPs), not memory bandwidth.

Why decode is memory-bound. During decoding, one token is generated per step. The attention computation is $O(L)$ where $L$ is the total sequence length so far (prompt + generated tokens) — linear in context length, but with a small constant (one query). The model must read the entire KV cache from HBM at every step, which is an $O(L \cdot d_{\text{cache}})$ memory operation. As $L$ grows, the bytes read per FLOP performed increases, and eventually the GPU's memory bandwidth becomes the limiting factor — the compute units idle waiting for data.

The TPLA (PD-sep.) strategy. During prefill, the system uses the reparameterized MLA form — the orthogonal transform $U$ has been absorbed into the weights ($W^{DKV}_{\mathrm{new}} = W^{DKV} U$, $W^{UKV}_{\mathrm{new}} = U^\top W_\gamma W^{UKV}$), but neither RMSNorm nor softmax is sliced. The computation is mathematically identical to the original MLA (modulo numerical precision effects from the transform), and tensor parallelism is applied along the head dimension — each GPU computes its assigned heads and the full latent cache is replicated. This preserves accuracy for the prefill tokens and leverages head-level TP to distribute the compute-bound workload.

The resulting KV cache is stored in the reparameterized latent space — that is, the cached vectors are $\mathbf{c} U$ rather than the original $\mathbf{c}$. This is critical because the TPLA decoding phase expects the cache in this transformed space.

During decoding, the system switches to TPLA mode: the latent cache (already stored as $\mathbf{c} U$) is split along the dimension axis across GPUs, each GPU computes local RMSNorm (with PCA correction factors) and local softmax (with variance-based scaling), and the partial outputs are all-reduced. The RoPE positional keys $\mathbf{k}^{\mathrm{PE}}$ are replicated across all GPUs since they are not latent-compressed and are needed for the full positional attention scores.

Why this helps accuracy. The prefill phase processes the entire prompt — potentially thousands of tokens — and its outputs (the KV cache for those tokens) influence all subsequent decoding steps. Using MLA without slicing for prefill means zero approximation error is introduced for the prompt tokens. The decoding phase introduces approximation error only for the newly generated tokens (typically much fewer than the prompt length), and even there, the error is in the attention computation over the cache (which is reparameterized but unsliced for prefill tokens, only dimension-split across GPUs). The paper reports that TPLA (PD-sep.) with no fine-tuning achieves near-original MLA accuracy on commonsense benchmarks and only a 2.15% average drop on LongBench for DeepSeek-V3, demonstrating that this separation largely sidesteps the conversion error.

KV cache reuse. An important implementation detail: the prefill KV cache is stored in the reparameterized space ($\mathbf{c} U$), so when decoding begins, the cache is already in the correct basis for TPLA's dimension-split attention. No conversion or re-encoding of the cache is needed at the phase boundary — the orthogonal transform having been absorbed into $W^{DKV}$ means the prefill computation naturally produces activations in the $U$-transformed space.

The TTFT (Time to First Token) benefit. Figure 4 shows that TPLA (sep.) achieves 1.4× faster prefill (lower TTFT) than vanilla TPLA for both DeepSeek-V3 and Kimi-K2 at 1K prompt length. This is because TPLA (sep.) uses head-level TP during prefill (each GPU computes $h_q/2$ heads rather than $h_q$), reducing per-device compute in the compute-bound phase. Vanilla TPLA, which keeps all heads on each device during prefill (since it also slices the latent dimension), lacks this head-level parallelism and thus is slower for prefill. The separation therefore improves both accuracy (no slicing error in prefill) and prefill latency (head TP reduces compute per device).

Complexity Analysis: Where the Speedup Comes From

The paper provides a rigorous FLOPs and memory traffic analysis (Section 4.5) comparing MLA-TP and TPLA-TP for the TP=2 case, explaining the throughput measurements in Figure 3.

Attention computation complexity. For the main attention operation (content-based, non-positional), with a query of length $L_q$ and KV cache of length $S_{KV}$:

• MLA-TP: Each GPU handles $h_q/2$ heads, each with $4d_h$ latent dimensions. Complexity: $\mathcal{O}(L_q \times S_{KV} \times \frac{h_q}{2} \times 4d_h \times 2)$, where the factor of 2 accounts for the two GPUs (though they operate in parallel, so wall-clock time divides this by 2).
• TPLA: Each GPU handles $h_q$ heads, each with $2d_h$ latent dimensions. Complexity: $\mathcal{O}(L_q \times S_{KV} \times h_q \times 2d_h \times 2)$.

These are arithmetically equivalent: $\frac{h_q}{2} \times 4d_h = h_q \times 2d_h = 2 h_q d_h$. The total number of multiply-add operations in the attention score computation is identical. The paper notes: "These two complexities are arithmetically equivalent."

Where TPLA has additional compute. The equivalence breaks down in two places:

1. Query projection: TPLA replicates the full $Q$ tensor on each device (all $h_q$ heads), while MLA-TP splits heads so each device computes only $h_q/2$ heads' queries. The query projection cost increases by approximately $\mathcal{O}(L_q \times D \times 2h_q \times d_h)$ for TPLA relative to MLA-TP.

2. RoPE computation: TPLA effectively doubles the number of heads (due to the GLA equivalence with duplication), so the RoPE positional attention scores $\mathbf{q}^{\mathrm{PE}} (\mathbf{k}^{\mathrm{PE}})^\top$ are computed for $h_q$ heads on each device rather than $h_q/2$. This adds minor overhead since the RoPE dimension $d_r$ is small (typically 64 per head).

3. All-reduce communication: TPLA requires an all-reduce of the output tensors $\tilde{O}_{\mathbf{0}}$ and $\tilde{O}_{\mathbf{1}}$ (each of size $B \times 1 \times D$), while MLA-TP typically uses an all-reduce on the output projection results as well (since head outputs need to be combined). The communication volume is similar.

Where TPLA saves memory bandwidth (the key advantage). The per-device KV cache size is:

• MLA-TP: Each GPU stores the full latent cache of $4d_h + d_r$ dimensions per token (e.g., 576 for DeepSeek-V3).
• TPLA: Each GPU stores $2d_h + d_r$ dimensions per token (e.g., 320 for DeepSeek-V3: $2 \times 128 = 256$ for the latent shard, plus $64$ for RoPE).

At a 32K context length, the per-device KV cache for DeepSeek-V3 (BF16, so 2 bytes per element) is $32000 \times 576 \times 2 \approx 36.9$ MB for MLA-TP versus $32000 \times 320 \times 2 \approx 20.5$ MB for TPLA — a 1.8× reduction. Since decoding is memory-bandwidth-bound, the throughput improvement is approximately proportional to the memory traffic reduction. The measured speedups of 1.79× and 1.93× (Figure 3) match this analysis closely, confirming that memory bandwidth savings dominate any additional compute overhead.

Why the overhead is acceptable. As context length grows, the attention computation (which is identical between MLA-TP and TPLA in total FLOPs) dominates the total cost. The additional query projection and RoPE overhead are constant per decoding step (independent of context length), while the memory bandwidth savings scale linearly with context length. At short contexts, the overhead might negate the benefit, but at the 32K+ context lengths where KV cache memory pressure matters most, the bandwidth savings far outweigh the extra FLOPs. Figure 3 confirms this: TPLA's throughput advantage grows with context length, reaching the ~2× asymptote at 32K for both models.

The $W^{VO}$ projection cost. After the attention-weighted sum $O_{\mathbf{0}}$ is computed (shape $B \times 1 \times h_q \times 2d_h$), it is projected through $W_{\mathbf{0}}^{VO} \in \mathbb{R}^{(h_q \cdot 2d_h) \times D}$. This projection has the same complexity as in MLA-TP where each GPU projects its $h_q/2$ heads through a $(h_q/2 \cdot 4d_h) \times D$ matrix — both are $\mathcal{O}(h_q \cdot 2d_h \cdot D)$ per GPU. The paper notes this equivalence: "Similarly, the $\tilde{O}$ computations are also equivalent."

In summary, TPLA's speedup comes from reducing the per-device KV cache size (and thus the memory bandwidth demand during decoding) by approximately a factor of $4d_h / (2d_h + d_r)$, at the cost of a small increase in query computation and the introduction of controlled approximation errors in RMSNorm and softmax that are mitigated through orthogonal reparameterization and prefill–decode separation.

4. Key Insights and Innovations

Innovation 1: Identifying the Structural Incompatibility Between MLA Compression and Tensor Parallelism as a First-Class Problem

The paper's most distinctive conceptual contribution is not a new architecture or training method, but rather a diagnosis: it identifies that Multi-Head Latent Attention (MLA) and tensor parallelism (TP) are structurally incompatible in a way that previous literature had not articulated. Prior work treated MLA's KV cache compression and TP's memory distribution as independently desirable techniques that could be combined straightforwardly — after all, standard multi-head attention partitions heads across GPUs naturally, so why would MLA be any different? The paper shows exactly why: MLA's compression works by merging per-head representations into a shared low-rank bottleneck ($\mathbf{c}^{\mathrm{KV}}$), and this very merging eliminates the per-head independence that TP normally exploits for memory distribution. Heads can be split across GPUs, but the shared latent cache cannot — it must be fully replicated on every device.

This diagnosis is significant because it reframes MLA's efficiency claim with a sharp boundary condition. MLA's per-token KV cache of 576 dimensions (for DeepSeek-V3) is genuinely smaller than GQA's per-token cache of 2048 dimensions on a single GPU — but under TP=4, MLA's per-device cache remains 576 while GQA's drops to 512. The paper makes this quantitative comparison explicit (Section 1, the LLaMA-3-70B vs. DeepSeek-V3 example), establishing that MLA's compression advantage over GQA reverses under sufficient tensor parallelism. This is not an implementation shortcoming — it is a consequence of the compression mechanism itself, and it means that the very technique designed to reduce memory becomes a memory liability at scale. The field had not articulated this tension clearly before; prior work on GLA (Zadouri et al., 2025) implicitly recognized the problem by proposing a new architecture, but did not frame it as a fundamental structural clash.

By naming and analyzing this incompatibility, the paper creates a new design axis for attention mechanisms: compressibility under parallelism. An attention mechanism should be evaluated not just on its single-GPU memory footprint, but on how that footprint scales with the tensor parallelism degree. This is a conceptual reframing that applies beyond MLA — any KV cache compression method that merges representations across heads will face the same tension with head-wise TP, and future designs should account for this from the start.

Innovation 2: Reparameterizing to Control Approximation Error in Decomposed Attention, Enabling Training-Free Conversion

The second major conceptual contribution is the reparameterization-as-error-control framing. TPLA decomposes the global attention computation into per-shard computations that are not mathematically equivalent due to the softmax non-linearity. Rather than accepting whatever accuracy loss this decomposition introduces (which would require retraining to recover), the paper asks: can we pre-transform the model's weights so that the decomposition error becomes small? The answer is yes, and the mechanism is inserting an orthogonal transform into the weight matrices via absorption.

What makes this distinctive is the inversion of the typical training-inference relationship. Normally, a model is trained, and inference optimizations (quantization, pruning, KV cache eviction) introduce degradation that is either tolerated or recovered through additional training. TPLA inverts this: it shows that through a purely algebraic reparameterization (absorb $U$ into $W^{DKV}$, absorb $U^\top$ and $W_\gamma$ into $W^{UKV}$), the model can be transformed into a different but mathematically equivalent form that, when subsequently decomposed, happens to have much smaller approximation error. The model's output under full-precision MLA is unchanged by the reparameterization alone — the error only appears when the latent dimension is split across GPUs and RMSNorm/softmax are computed locally. The quality of the reparameterization determines how much the split version deviates from the original.

This framing is more sophisticated than a simple "approximate and fine-tune" approach. It recognizes that not all weight-space transformations are equal under decomposition — some bases (the original learned one) amplify the decomposition error, while others (PCA-rotated, Hadamard-rotated) suppress it. The paper's Conditions 1 and 2 formalize what properties a good basis should have: norm proportionality for RMSNorm slicing and score proportionality for softmax slicing. These conditions are not derived from first principles of attention — they are engineering heuristics that identify what statistical properties of the transformed activations minimize the per-shard approximation error. The fact that these conditions can be satisfied through simple data-driven PCA on a small calibration corpus (WikiText-2, in the experiments) makes the approach practical without requiring any modification to the training pipeline.

The significance extends beyond TPLA itself. The reparameterization-as-error-control pattern — identify an inference-time decomposition that would be efficient if only the weights were in a better basis, then absorb an orthogonal transform to achieve that basis without changing the model's function — is generalizable. It could apply to other attention variants, other normalization schemes, or other forms of tensor parallelism where a global non-linearity must be approximated by local computations. The paper does not make this generalization claim explicitly, but the framework is transferable.

Innovation 3: Prefill–Decode Separation as a Phase-Aware Strategy That Sidesteps Rather Than Solves the Approximation Problem

The third conceptual insight is the prefill–decode separation (PD-sep.) strategy, which takes a pragmatic approach to the approximation error problem: rather than trying to make the decomposition error zero everywhere (impossible without retraining), confine it to the phase where it matters least, and eliminate it from the phase where it matters most.

This is not just an engineering optimization — it reflects a deeper understanding of the error propagation structure in autoregressive inference. During prefill, the model processes thousands of prompt tokens and produces KV cache entries that will be attended to by every subsequent decoding step. Approximation errors introduced during prefill compound: a slightly wrong KV cache entry affects all future tokens' attention distributions, creating a cascading degradation that grows with sequence length. During decoding, by contrast, only one token is generated per step, and its attention computation affects only that step's output — errors do not accumulate in the KV cache (the newly generated token's KV entry is its own; errors in its computation do not affect the stored representation of previous tokens).

PD-sep exploits this asymmetry. By using unsliced MLA during prefill (reparameterized but not decomposed), the KV cache is stored with zero slicing error for the vast majority of tokens that will ever be in the cache. When decoding switches to TPLA with sliced RMSNorm and softmax, the approximation error only affects the attention computation over this high-quality cache — the cache itself is not degraded by the slicing. The result, shown in Table 2, is that the training-free TPLA (PD-sep.) on DeepSeek-V3 loses only 2.15% average accuracy on LongBench, a benchmark specifically designed to stress long-context understanding where cumulative errors would be most visible.

This strategy is conceptually distinct from the typical "train to recover accuracy" approach that dominates inference optimization papers. It represents a structural solution rather than a data-driven one: by architecting the inference pipeline to match the error tolerance of each phase (zero-tolerance for prefill, moderate-tolerance for decode), the paper achieves near-original accuracy without any training. The fact that TPLA (PD-sep.) outperforms the aligned variant (which used 100M tokens of fine-tuning) on DeepSeek-V2-Lite LongBench (Table 2) is particularly telling — it suggests that targeted structural correctness can beat blunt-force fine-tuning for this type of decomposition error. The PD-sep. strategy is likely applicable to other inference optimizations where an approximation error can be phase-gated, making it a reusable architectural pattern beyond TPLA.

Innovation 4: TPLA as GLA with Doubled Heads — An Algebraic Equivalence That Unlocks Kernel Compatibility

The fourth contribution is the algebraic equivalence established in Section 4.4: TPLA with its latent-dimension splitting, full-head replication, and all-reduce output combination is mathematically equivalent to a Grouped Latent Attention (GLA) system with twice as many attention heads (for $g=2$ groups) and a specific block structure in the query tensor. This equivalence is not just a theoretical curiosity — it has immediate practical consequences.

First, it means TPLA can directly leverage optimized attention kernels built for GLA, particularly FlashAttention-3, without requiring new low-level implementations. The paper states this explicitly as a practical advantage: "state-of-the-art attention optimizations (e.g., FlashAttention-3) can be applied to TPLA without substantial changes to the underlying framework." This is significant because high-performance attention kernels require enormous engineering effort to develop and tune — inheriting GLA's kernel compatibility dramatically reduces the barrier to deploying TPLA in production systems.

Second, the equivalence clarifies why TPLA preserves representational capacity while GLA loses it. In GLA, each group of heads sees only its assigned latent shard — the off-diagonal query blocks $Q_{\mathbf{0,1}}$ and $Q_{\mathbf{1,0}}$ are simply discarded, so each head's computation uses only half the latent information. In TPLA, the head duplication means that what GLA discards as off-diagonal blocks, TPLA retains as on-diagonal blocks of the duplicated heads — every latent shard is attended to by a complete set of heads. The representational cost of GLA is not inherent to latent-dimension partitioning; it is an artifact of how the heads are grouped. TPLA shows that latent-dimension partitioning and full-head access are not in tension — they can coexist through the head-duplication trick.

Third, the equivalence provides a conceptual bridge for future work. If TPLA is GLA with doubled heads, then improvements to GLA's sharding patterns, communication strategies, or kernel implementations automatically benefit TPLA. Conversely, TPLA's reparameterization techniques (Hadamard/PCA transforms to control softmax slicing error) could potentially be applied to GLA models to improve their accuracy without architectural changes, since both architectures face the same decomposed-softmax approximation problem.

Innovation 5: Verifier-Free Conversion — Preserving Accuracy Without a Learned Scoring Function

A more subtle but practically crucial innovation is what TPLA does not require: any learned verifier, reward model, or teacher-student distillation to validate or recover the converted model's outputs. This contrasts sharply with the dominant paradigm in the inference optimization literature, where compression, quantization, or architectural changes typically require either (a) a separate trained model to judge output quality, (b) access to ground-truth labels for fine-tuning, or (c) a teacher model for distillation.

TPLA's conversion pipeline uses only a small calibration corpus (WikiText-2, publicly available) to compute PCA eigenvectors — no labels, no teacher model, no reinforcement learning, no iterative self-improvement. The orthogonal transform is derived purely from activation statistics. The optional alignment step uses 100M tokens of SmolLM-Corpus with a simple MSE loss matching layer-wise features to the original MLA model — again, no external verifier or task-specific supervision. This means TPLA can be applied to any MLA-pretrained model (DeepSeek series, Kimi-K2, TransMLA-converted models) without access to the original training data, task labels, or even the task definition. The conversion is self-contained — it depends only on the model's own weights and activations.

This property is particularly important for large-scale deployment. When a new MLA-based model is released, practitioners can apply TPLA and immediately benefit from the throughput improvements without setting up an evaluation pipeline to verify that conversion hasn't broken task-specific performance. The paper's evidence that the training-free PD-sep. variant achieves near-original accuracy across diverse benchmarks (commonsense reasoning in Table 1, LongBench in Table 2) provides confidence that the conversion is robust across task types, not just calibrated to a specific evaluation. In an ecosystem where model releases are frequent and deployment efficiency is critical, this zero-verification property is a substantial practical advantage over methods that require task-specific accuracy validation after conversion.

The contrast with GLA is instructive here: GLA cannot be applied to existing MLA checkpoints at all without catastrophic degradation (WikiText-2 perplexity jumps from 6.31 to 2212, Table 1), making it strictly a from-scratch-training proposal. TPLA's ability to convert existing checkpoints with bounded, predictable accuracy loss makes it immediately deployable on the current generation of MLA-based models — DeepSeek-V2, DeepSeek-V3, and Kimi-K2 — which collectively represent some of the most capable open-weight LLMs. This practical immediacy, enabled by the verifier-free conversion design, is arguably more important for real-world impact than the architectural novelty itself.

5. Experimental Analysis

Evaluation Methodology

• Dataset. The paper evaluates on two categories of benchmarks. For short-text commonsense reasoning (Table 1), the evaluation uses WikiText-2 (perplexity) and six commonsense benchmarks through the LightEval framework: MMLU, ARC (Easy and Challenge), PIQA, HellaSwag, OpenBookQA (OBQA), and WinoGrande (WG). For long-context understanding (Table 2), the evaluation uses LongBench, a bilingual (English/Chinese) multi-task benchmark comprising 21 tasks across six categories including question answering, summarization, and few-shot learning. Due to GPU memory constraints, maximum input context length is set to 31,500 tokens for DeepSeek-V2-Lite and 127,500 tokens for DeepSeek-V3, with output lengths matching the original paper's settings.

• Base model(s). The paper evaluates on DeepSeek-V2-Lite (converted from MLA via TransMLA), DeepSeek-V3-0324 (685B parameters, Mixture-of-Experts), and Kimi-K2-Base (1T parameters). LLaMA-2-7B is also used as an initial testbed after conversion through TransMLA (which maps MHA/GQA to MLA with 64 RoPE dimensions and 512 NoPE dimensions, corresponding to a 92.97% pruning ratio), followed by TPLA conversion. For throughput and latency measurements, the MoE routing is removed from DeepSeek-V3 and Kimi-K2 to isolate attention-speed effects. All models are converted to BF16 for inference. DeepSeek-V2-Lite serves as the primary testbed for accuracy ablations because its smaller scale enables faster experimentation, while DeepSeek-V3 and Kimi-K2 demonstrate the approach at production scale.

• Metrics. Three categories of metrics are reported. Accuracy metrics: WikiText-2 perplexity (lower is better) for language modeling quality; accuracy (%) on each of the six commonsense benchmarks; and LongBench average accuracy across its 21 tasks. Throughput: decoding throughput measured in tokens per second under maximum batch size at each context length, comparing MLA and TPLA on two GPUs (Figure 3). Latency: Time to First Token (TTFT) measured in milliseconds for the prefill phase, comparing TPLA and TPLA (sep.) on two GPUs (Figure 4).

• Baselines. The paper compares against several baselines and variants. Original MLA: the unmodified pretrained model using standard Multi-Head Latent Attention, serving as the accuracy upper bound. GLA (Grouped Latent Attention): direct conversion from MLA to GLA following Section 3.2, partitioning attention heads into two groups with each group assigned half the latent dimension. This represents the prior approach most similar in objective to TPLA. MLA-TP: standard MLA under tensor parallelism with heads split across devices but the full latent cache replicated, serving as the throughput baseline for decoding speed comparisons. TPLA (vanilla): the proposed method without prefill–decode separation — RMSNorm and softmax are sliced during both prefill and decode. TPLA (align): TPLA with lightweight alignment using the SmolLM-Corpus (100M tokens, following TransMLA settings with batch size 32, learning rate 2e-5, warmup ratio 0.03, cosine scheduler, max sequence length 4096). TPLA (PD-sep.): TPLA with prefill–decode separation, using reparameterized MLA without slicing during prefill and TPLA with slicing during decoding, training-free.

• Generation budget / compute accounting. The paper does not define a "generation budget" in terms of sample counts, since this is a systems paper focused on inference architecture rather than output sampling strategies. Instead, compute is measured through three system-level metrics. For complexity analysis (Section 4.5), FLOPs are counted in terms of the asymptotic operations for attention score computation, query projection, and output projection under both MLA-TP and TPLA-TP, with per-device KV cache size used as the proxy for memory bandwidth demand. For throughput experiments (Section 5.4.1, Figure 3), the metric is tokens per second measured on two GPUs with maximum batch size configured at each context length, using FlashAttention-3 for all implementations. For latency experiments (Section 5.4.2, Figure 4), Time to First Token (TTFT) is measured on two GPUs for the prefill phase. Memory traffic reduction is quantified through the per-device KV cache dimensionality: MLA stores 576 dimensions per token per device (4dh + dr = 512 + 64), while TPLA stores 320 dimensions per token per device (2dh + dr = 256 + 64), yielding approximately a 1.8x reduction that directly translates to throughput improvement since decoding is memory-bandwidth-bound.

• Cross-validation / statistical protocol. No cross-validation protocol is reported, since this is not a strategy-selection paper. The reparameterization uses WikiText-2 as a calibration set for PCA eigenvector computation. The alignment procedure uses SmolLM-Corpus. Evaluation is performed on the standard test sets of each benchmark. The throughput and latency measurements use direct GPU profiling rather than statistical estimation. The paper reports results on the full test sets without mentioning train/validation splits for hyperparameter tuning of the conversion process.

Main Quantitative Results

Commonsense Reasoning and Language Modeling Accuracy (Table 1)

Table 1 reports WikiText-2 perplexity and six commonsense benchmark accuracies for DeepSeek-V2-Lite under different conversion approaches. The headline finding is that TPLA preserves MLA's representational capacity with minimal degradation, while GLA causes catastrophic collapse.

GLA conversion destroys performance. On WikiText-2, perplexity explodes from 6.31 (MLA) to 2212 (GLA) — a greater than 350× degradation. The commonsense benchmarks show corresponding collapse: MMLU drops from 55.0% (MLA) to 24.04% (GLA); ARC-Easy drops from 77.19% to 52.33%; HellaSwag drops from 64.69% to 32.03%. This confirms the paper's claim that "discarding half of each head's KV cache causes severe performance degradation" and that GLA's representational restriction (each head accessing only half the latent dimension) is fundamentally incompatible with MLA-pretrained weights.

Training-free TPLA with PD separation nearly matches MLA. TPLA (PD-sep.), which requires no training, achieves WikiText-2 perplexity of 6.36 — only 0.05 higher than the original MLA's 6.31. On the commonsense benchmarks: MMLU drops from 55.0% to 54.01% (−0.99 points); ARC-Easy from 77.19% to 76.77% (−0.42 points); ARC-Challenge from 44.97% to 44.20% (−0.77 points); PIQA from 77.86% to 78.02% (+0.16 points, actually slightly better); HellaSwag from 64.69% to 64.27% (−0.42 points); OBQA from 42.80% to 41.80% (−1.00 points); WinoGrande from 68.67% to 68.90% (+0.23 points). No benchmark shows more than a 1.0 percentage point drop. This supports the paper's claim that prefill–decode separation "achieves performance close to the original model without any training."

PCA-based reparameterization with alignment recovers performance. TPLA (align), which applies PCA-based reparameterization and then fine-tunes on 100M tokens of SmolLM-Corpus, achieves WikiText-2 perplexity of 6.68 — slightly higher (worse) than PD-sep. but substantially better than the unaligned TPLA (7.24). On commonsense benchmarks, TPLA (align) generally matches or slightly trails PD-sep.: MMLU 54.13% vs 54.01%, ARC-Easy 76.22% vs 76.77%, ARC-Challenge 44.54% vs 44.20%, PIQA 77.48% vs 78.02%, HellaSwag 63.86% vs 64.27%, OBQA 41.20% vs 41.80%, WinoGrande 69.14% vs 68.90%. The recovery is effective but PD-sep.'s training-free advantage is notable — the alignment step provides only marginal benefit over the structurally sound PD-sep. approach.

Vanilla TPLA without PD separation shows moderate degradation. TPLA (vanilla) uses PCA reparameterization but slices RMSNorm and softmax during both prefill and decoding, without alignment. WikiText-2 perplexity rises to 7.24 (from 6.31), and commonsense benchmarks show drops of approximately 1–4 percentage points: MMLU 52.18% (−2.82), ARC-Easy 74.33% (−2.86), ARC-Challenge 41.47% (−3.50), PIQA 75.35% (−2.51), HellaSwag 61.33% (−3.36), OBQA 40.60% (−2.20), WinoGrande 66.93% (−1.74). This degradation, while not catastrophic like GLA's, demonstrates that slicing RMSNorm and softmax during prefill introduces meaningful approximation error that PD separation successfully avoids.

TransMLA bridge enables TPLA for GQA/MHA models. For LLaMA-2-7B, which uses GQA, the conversion pipeline is: TransMLA converts GQA to MLA → fine-tune to recover performance → TPLA conversion. The paper reports that this bridge pathway enables TPLA to be "applied to pretrained models that originally use MLA, GQA, or MHA," with the TransMLA checkpoint released publicly. Specific accuracy numbers for LLaMA-2-7B TPLA are not reported in Table 1, but the demonstration of the conversion pathway is itself part of the claim that TPLA is broadly applicable.

Long-Context Understanding: LongBench Results (Table 2)

Table 2 reports LongBench accuracy for DeepSeek-V2-Lite and DeepSeek-V3 under different conversion approaches at the maximum feasible context lengths (31,500 and 127,500 tokens respectively). The headline finding is that TPLA (PD-sep.) maintains strong long-context performance with only modest degradation, while alignment on short-text corpora does not reliably transfer to long-context tasks.

DeepSeek-V2-Lite. The original MLA achieves 41.82% average accuracy. TPLA (vanilla) drops to 38.95% (−2.87 points). TPLA (align), which uses short-text SmolLM-Corpus for fine-tuning, achieves 38.53% (−3.29 points) — actually worse than vanilla TPLA, suggesting that alignment on concatenated short texts does not help and may even harm long-context capabilities. TPLA (PD-sep.) achieves 40.02% (−1.80 points), the best among TPLA variants and notably surpassing the aligned variant. The paper interprets this as evidence that "slicing errors in RMSNorm and softmax accumulate with sequence length" (affecting vanilla TPLA more on LongBench than on short commonsense tasks), and that "the alignment corpus is formed by concatenating short texts" (limiting its effectiveness for genuine long-context understanding). The PD-sep. variant's structural avoidance of prefill slicing errors proves more effective than post-hoc short-text fine-tuning for preserving long-context accuracy.

DeepSeek-V3. The original MLA achieves 56.59% average accuracy (averaged across all LongBench tasks; the paper reports the average of the category scores). TPLA (vanilla) drops to 52.51% (−4.08 points). TPLA (align) achieves 53.48% (−3.11 points), showing some recovery but leaving a substantial gap. TPLA (PD-sep.) achieves 54.44% (−2.15 points), the best among TPLA variants. The paper highlights this as a key finding: "On DeepSeek-V3 the model retains strong long-form reasoning with only a modest average drop of 2.15%." The 2.15% degradation is framed as acceptable for deployment, especially since it is "likely recoverable with a small amount of additional training" on long-context data specifically (unlike the short-text alignment attempted here).

Category-level analysis. The paper does not report per-category LongBench scores in Table 2, only the overall average. This is a limitation — certain categories (e.g., summarization vs. few-shot learning) may have different sensitivity to the softmax slicing approximation error, and category-level results would reveal whether the degradation is uniform or concentrated in specific task types. Without this breakdown, the average score obscures potentially important variance.

Decoding Throughput Speedup (Figure 3)

Figure 3 plots decoding throughput (tokens per second) against context length for MLA-TP and TPLA on two GPUs for DeepSeek-V3-0324 and Kimi-K2-Base (both with MoE removed to isolate attention effects). The headline numbers announced in the abstract are confirmed.

DeepSeek-V3-0324 (Figure 3a). At 4K context length, TPLA achieves approximately 2× the throughput of MLA — the maximum gain observed. As context length increases, the absolute throughput of both methods decreases (more cache to read per step), but TPLA's advantage narrows slightly while remaining substantial. At 32K context length, TPLA achieves 1.79× the throughput of MLA. The paper reports that the "maximum batch size at each context length" is configured, meaning the comparison accounts for batching efficiency, not just single-request latency.

Kimi-K2-Base (Figure 3b). The pattern is similar. At 4K context length, TPLA achieves approximately 2× throughput. At 32K context length, TPLA achieves 1.93× the throughput of MLA — a slightly larger gain than DeepSeek-V3, likely due to Kimi-K2's different head configuration (32 heads vs. DeepSeek-V3's 64 heads, affecting the ratio of attention computation to memory traffic).

Why the speedup approaches but doesn't exceed 2×. The theoretical maximum memory traffic reduction is the ratio of per-device KV cache sizes: for DeepSeek-V3, 576/320 = 1.8×; for Kimi-K2, the paper implies a similar ratio though exact KV cache dimensions are not provided. The measured speedups of 1.79× and 1.93× match this theoretical bound closely, confirming that decoding is indeed memory-bandwidth-bound under these configurations and that TPLA's additional computational overhead (query replication, RoPE duplication) does not meaningfully reduce the net benefit at long contexts. The sub-2× speedup reflects that (a) not all memory traffic is KV cache reads (model weights also need to be streamed), and (b) TPLA introduces some additional computation per step.

Note on TP>2. The paper states that for TP>2, they "further split heads in addition to halving the latent dimension." With TP=4 on Kimi-K2-TPLA, each device gets 32 heads × 2 with a 320-dimensional latent per head. In this configuration, "the per-device compute halves, while memory traffic matches TP=2; decoding remains memory-bound, so the speedup is similar to TP=2." Measurements are therefore reported only for two GPUs, since higher TP degrees do not change the per-device KV cache size under TPLA's current scheme (the latent dimension is only split once, into two halves, regardless of TP degree).

Prefill Latency Reduction (Figure 4)

Figure 4 plots Time to First Token (TTFT) against prompt length for TPLA and TPLA (sep.) on two GPUs for DeepSeek-V3-0324 and Kimi-K2-Base (MoE removed).

DeepSeek-V3-0324 (Figure 4a). At a 1K prompt length, TPLA (sep.) achieves 1.4× faster TTFT than vanilla TPLA. The gap widens slightly at longer prompt lengths as the compute-bound prefill phase becomes increasingly dominated by the attention computation (which scales quadratically with prompt length), where TPLA (sep.)'s head-level TP provides greater benefit. The paper describes this as "essentially a 'free lunch'" because the latency improvement comes from the structurally more appropriate TP strategy for the compute-bound phase without introducing any accuracy degradation (since PD-sep. uses unsliced MLA during prefill).

Kimi-K2-Base (Figure 4b). The same 1.4× speedup is observed at 1K prompt length. The consistency across both model architectures suggests the benefit is robust and primarily determined by the ratio of per-device head computation between the two strategies (TPLA processes all heads per device during prefill, while TPLA (sep.) processes half).

Why TPLA (sep.) is faster for prefill. During prefill, TPLA keeps all $h_q$ heads on each device (since it partitions the latent dimension, not heads). TPLA (sep.) uses unsliced MLA with head-level TP — each GPU computes only $h_q/2$ heads, reducing per-device FLOPs proportionally. In the compute-bound prefill phase, this compute reduction directly translates to lower latency. The 1.4× measurement (rather than the theoretical 2× from halving heads) reflects that not all prefill computation is in the attention module — the MLP layers, embedding lookup, and other operations are not affected by the TP strategy for attention.

Ablation Studies and Robustness Checks

MLA → GLA direct conversion (Table 1, Figure 2). This experiment tests what happens when an MLA-pretrained model is converted to GLA by partitioning attention heads and latent dimensions as described in Section 3.2, without any retraining or reparameterization. The result is catastrophic: WikiText-2 perplexity explodes from 6.31 (MLA) to 2212 (GLA), and all commonsense benchmarks show massive degradation (MMLU drops from 55.0% to 24.04%, a 31-point collapse). This serves as a negative control demonstrating that GLA's structural restriction — each head accessing only half the latent dimension — is fundamentally incompatible with weights trained under the full-dimensional MLA regime. The paper uses this result to motivate TPLA's design: the problem is not latent-dimension partitioning per se, but the loss of per-head access to the full latent information.

Prefill–decode separation on/off (Table 1, Table 2). Comparing TPLA (vanilla) against TPLA (PD-sep.) isolates the effect of slicing RMSNorm and softmax during prefill versus only during decoding. On WikiText-2, PD-sep. achieves 6.36 perplexity vs. 7.24 for vanilla TPLA — the prefill slicing introduces approximately 13% of the total degradation from MLA to vanilla TPLA (since 7.24 − 6.31 = 0.93 total degradation, and 7.24 − 6.36 = 0.88 degradation attributed to prefill slicing). On LongBench for DeepSeek-V3, PD-sep. achieves 54.44% vs. 52.51% for vanilla, a 1.93-point improvement. The paper interprets this as evidence that "slicing errors in RMSNorm and softmax accumulate with sequence length," making the prefill phase particularly sensitive since its errors propagate through the entire KV cache. The finding that PD-sep. outperforms the aligned variant (which uses 100M tokens of fine-tuning) on DeepSeek-V2-Lite LongBench (40.02% vs. 38.53%) is non-obvious and suggests that structural avoidance of prefill error is more effective than post-hoc short-text fine-tuning for long-context tasks.

Reparameterization method: Hadamard vs. PCA vs. Original (Figure 2). This ablation (reported without PD separation to make the method-induced loss more visible) decomposes TPLA's approximation error along two axes: which component is parallelized (RMSNorm only, softmax only, or both) and which reparameterization strategy is used (original weight splitting, Hadamard transform, or PCA). The key findings are fourfold.

First, error ordering: slicing RMSNorm alone incurs the least performance loss, slicing softmax alone is worse, and slicing both is worst — confirming that softmax decomposition is the primary source of approximation error.

Second, Hadamard for RMSNorm slicing: the Hadamard-based method "balances the norm computation across devices effectively, leading to performance comparable to the original MLA model on multiple tasks" (Figure 2, "TPLA (norm only)" with Hadamard texture). This validates Condition 1 — the Hadamard transform's uniform energy distribution makes per-shard RMS values accurately approximate the global RMS.

Third, Hadamard for softmax slicing fails: "the Hadamard-based method fails to improve softmax accuracy" (Figure 2, "TPLA (softmax only)" with Hadamard texture shows performance barely above the "original" splitting baseline). The paper hypothesizes this is because "the exponential nature of softmax makes it more sensitive to imbalance," and while Hadamard achieves statistical balance in expectation, it produces large per-sample asymmetries in the attention scores due to the mixed-sign contributions (as demonstrated by the counterexample in Section 4.3.1). This is a genuinely non-obvious finding — a transform that perfectly balances norms can simultaneously fail to balance dot products because the signed contributions cancel globally but not locally.

Fourth, PCA dominates when both are sliced: "When both components are parallelized, the PCA-based reparameterization consistently achieves the best performance" (Figure 2, "TPLA" with PCA texture showing the smallest accuracy drop below the MLA baseline). This leads to the paper's adoption of PCA as the default reparameterization for all main experiments.

Alignment effectiveness (Table 1, Table 2). The alignment procedure involves two stages: first, matching layer-wise input/output features of TPLA to the original MLA model using 256 random samples of length 2,048 for 10 epochs with MSE loss and the Muon optimizer (learning rate 1e−6); second, aligning end-to-end model outputs using 100M tokens from SmolLM-Corpus following TransMLA settings (batch size 32, learning rate 2e−5, warmup ratio 0.03, cosine scheduler, max sequence length 4096). On commonsense benchmarks (Table 1), alignment recovers most of the vanilla TPLA degradation: TPLA (align) achieves MMLU 54.13% vs. vanilla TPLA's 52.18% (+1.95 points) and MLA's 55.0%. On LongBench for DeepSeek-V2-Lite (Table 2), alignment actually underperforms vanilla TPLA (38.53% vs. 38.95%), likely because "the alignment corpus is formed by concatenating short texts" and does not train the model to handle the cumulative softmax slicing errors that manifest at long context lengths. For DeepSeek-V3, alignment provides modest recovery (53.48% vs. 52.51% for vanilla TPLA) but still trails PD-sep. (54.44%). These results suggest that alignment is a partial fix at best, and its effectiveness depends on the match between the alignment data distribution and the target task distribution — a limitation the paper acknowledges only briefly.

Calibration dataset for PCA (implicit ablation). The paper uses WikiText-2 as the calibration dataset for PCA eigenvector computation and SmolLM-Corpus for the alignment procedure. No ablation is reported comparing different calibration datasets to assess how sensitive the PCA-derived eigenvectors and variance fractions ($\alpha$, $\beta$) are to the choice of calibration data. This is a significant gap: if PCA eigenvectors computed from WikiText-2 (English Wikipedia-style text) do not transfer well to other domains (code, math, multilingual text), the conversion quality could degrade in deployment. The strong LongBench results (which includes Chinese and diverse task types) provide indirect evidence of reasonable transfer, but a direct ablation would have strengthened this claim.

TPLA as GLA compatibility (Section 4.4, not a separate experiment). The paper establishes that TPLA is algebraically equivalent to GLA with doubled heads, which means "state-of-the-art attention optimizations (e.g., FlashAttention-3) can be applied to TPLA without substantial changes to the underlying framework." The throughput experiments in Figure 3 use FlashAttention-3 for both MLA and TPLA implementations, confirming this compatibility operationally. However, no ablation compares TPLA performance with and without FlashAttention-3 to quantify the kernel-level benefit specifically.

Effect of $g$ > 2 (not experimentally validated). The paper discusses extending TPLA to more than two latent-cache groups ($g > 2$) in Section 4.4 and notes that PCA "probably fails to maintain effectiveness" when $g > 2$ because the middle partitions would contain moderate-variance dimensions that are neither dominant enough to approximate the global computation nor negligible enough to ignore. No experiments with $g > 2$ are reported, making this an acknowledged but unexplored limitation. The paper suggests that Hadamard-based transforms may be more suitable for $g > 2$ if the softmax slicing problem can be solved, but this remains future work.

Critical Assessment

Does TPLA actually preserve MLA's representational capacity while enabling TP efficiency?

This is the paper's central claim. The evidence supports it, but with a nuance. What was demonstrated: Converting DeepSeek-V2-Lite and DeepSeek-V3 from MLA to TPLA (PD-sep.) without training preserves commonsense accuracy within approximately 1 percentage point across all six benchmarks (Table 1) and long-context accuracy within 2.15% on LongBench (Table 2). This is genuinely strong — TPLA is not just "less bad" than GLA; it achieves near-parity with the original MLA on a diverse set of tasks. What was not demonstrated: Whether TPLA actually preserves "representational capacity" in the theoretical sense (i.e., whether the attention mechanism can express the same class of functions). The paper's evidence is empirical rather than theoretical — there is no proof that the softmax decomposition error is bounded for all possible inputs under PCA reparameterization, only that it is small on the tested benchmarks. The "representational capacity" claim is inferred from the accuracy preservation, not directly established. This is a reasonable inference but should be understood as empirical validation rather than theoretical guarantee.

Are the throughput speedups (1.79×, 1.93×) genuine and attributable to TPLA's innovations?

The speedups are credible and well-measured, but two qualifications apply. First, the MoE routing is removed for the throughput measurements. While the paper states this is to "isolate attention-speed effects," it means the reported numbers apply to the dense attention component only; in a full MoE deployment, the speedup as a fraction of total inference time would be lower because the MoE computation (which is unaffected by TPLA) contributes to overall latency. The 1.79× and 1.93× figures should therefore be interpreted as speedups on the attention component, not end-to-end throughput improvements for the full MoE model. Second, the measurements are on two GPUs with TP=2. As the paper notes, for TP>2, the per-device KV cache size does not further decrease under TPLA's current scheme because the latent dimension is only split once (into two halves), and additional parallelism comes from splitting heads within each latent group. This means the memory bandwidth benefit does not scale with TP degree beyond 2 — a practical limitation for deployments requiring higher TP degrees.

Does TPLA (PD-sep.) truly require no training?

Yes, for the results reported in Tables 1 and 2, TPLA (PD-sep.) uses no training data and no gradient updates — it is purely a weight reparameterization plus a phase-dependent attention switching strategy. However, the PCA transform requires a calibration dataset (WikiText-2) to compute eigenvectors. While this is not "training" in the gradient sense, it does require running the model on a representative corpus to collect activations. The quality of the PCA transform depends on this calibration data, and the paper does not test sensitivity to the calibration data choice (e.g., would PCA computed from code corpora work for code generation tasks?). Additionally, the "no training" claim applies to the PD-sep. variant specifically; the TPLA (align) variant uses 100M tokens of fine-tuning, which the results show is less effective than PD-sep. for long-context tasks.

Does the paper adequately compare against alternative solutions?

The primary comparison is TPLA vs. standard MLA-TP and TPLA vs. GLA. These are the right baselines given the paper's framing. However, several comparisons are absent. No comparison against KV cache quantization: deploying MLA with int8 or int4 KV cache quantization would also reduce memory traffic, potentially achieving similar throughput gains without changing the attention structure. A FLOPs-matched or memory-matched comparison against quantized MLA would clarify whether TPLA's approach is complementary or redundant with quantization. No comparison against sequence parallelism: the paper mentions sequence parallelism in the related work but does not compare TPLA's throughput against, e.g., splitting the sequence across GPUs with MLA. No comparison against token eviction methods: techniques like H2O or SnapKV that prune low-importance tokens would also reduce effective KV cache size and could be applied on top of MLA. It is unclear whether TPLA's throughput benefit is competitive with or orthogonal to these approaches. No GLA from scratch: the paper compares against GLA only in the zero-shot conversion setting (MLA weights → GLA structure), where GLA catastrophically fails. This is appropriate for the paper's claim about conversion, but does not test whether a properly trained GLA model would match TPLA's accuracy-efficiency tradeoff. The paper acknowledges this indirectly by noting GLA "requires training from scratch" and leaving TPLA from-scratch training to future work, but a fair comparison would require training GLA and TPLA from scratch and comparing their accuracy-efficiency curves.

Are the accuracy results robust and comprehensive enough?

The benchmarks cover a reasonable range: language modeling (WikiText-2), commonsense reasoning (6 tasks), and long-context understanding (LongBench with 21 tasks). However, several gaps exist. No instruction-following or chat benchmarks: the models evaluated (DeepSeek-V2-Lite, DeepSeek-V3) are intended for chat and instruction-following use cases, but only academic accuracy benchmarks are reported. It is unknown whether TPLA conversion preserves instruction-following quality, safety alignment, or refusal behavior. No code or math benchmarks: these are key use cases for the DeepSeek models; their absence leaves open whether the softmax slicing approximation error disproportionately affects structured reasoning tasks. Single calibration dataset: all reparameterization uses WikiText-2 for PCA; robustness to distribution shift in the calibration data is not tested. No statistical significance or variance: results are reported as single accuracy numbers without confidence intervals, standard deviations, or multiple random seeds. Given the test sets have finite size (e.g., MMLU with ~14K questions, but per-category subsets are smaller), small differences (e.g., 54.01% vs. 54.13%) may not be statistically distinguishable. The paper treats numerical differences as meaningful without establishing what constitutes a significant difference.

What experiments would have strengthened the paper?

Several experiments are conspicuous by their absence. (1) A direct comparison on the full MoE models (with routing enabled) showing end-to-end throughput improvement, not just attention-component speedup. The paper removes MoE "to isolate attention-speed effects," but the practical value of TPLA depends on its end-to-end impact. (2) A sensitivity analysis varying the calibration dataset for PCA (e.g., WikiText-2 vs. C4 vs. code corpora vs. multilingual corpora) and measuring the resulting accuracy variance. This would establish whether the method requires careful calibration data selection or is robust to data choice. (3) A layer-wise analysis of where the softmax slicing error accumulates most. The paper treats all transformer layers uniformly, but different layers may have different activation statistics and different sensitivity to the approximation. (4) Scaling to $g>2$ groups with higher TP degrees, even if only to demonstrate the failure mode and quantify when PCA-based reparameterization breaks down. (5) Combination with KV cache quantization to demonstrate orthogonality — if TPLA + int8 quantization yields benefits multiplicative with either alone, the practical value proposition strengthens considerably. (6) A measurement of the actual memory bandwidth reduction using GPU profiling tools (e.g., NVIDIA Nsight) to directly confirm that the throughput improvement correlates with reduced HBM traffic, rather than inferring it from cache size ratios. (7) Evaluation on code generation benchmarks (HumanEval, MBPP) and mathematical reasoning (GSM8K, MATH) to test whether the softmax slicing error has domain-specific effects.

Do the PD-sep. results genuinely demonstrate that no training is needed, or is there a hidden cost?

The PD-sep. variant incurs no gradient-based training, but the paper does not account for the one-time calibration cost of computing PCA eigenvectors. For a model at DeepSeek-V3 scale (685B parameters), running forward passes on a calibration dataset (WikiText-2, ~2M tokens) to collect KV latent activations across all layers is computationally non-trivial, though far cheaper than fine-tuning. The paper does not report this calibration cost or compare it to the alignment cost. Additionally, the PCA transform must be recomputed if the model is fine-tuned on new data — the eigenvectors are specific to the activation statistics of the current weights. For a continuously trained or periodically updated model, this recalibration cost would recur. The paper frames PD-sep. as "training-free," which is accurate in the gradient sense, but the one-time calibration is a form of data-dependent adaptation that a truly zero-shot conversion (like int8 quantization with fixed scaling factors) does not require.

Summary assessment

The experimental evidence demonstrates that TPLA achieves its primary objective: converting MLA-pretrained models to a tensor-parallel-friendly form with bounded accuracy loss ($\leq$2.15% on long-context tasks, $\leq$1% on commonsense tasks) while delivering approximately 1.8–1.9× decoding throughput improvements at 32K context length. The GLA comparison establishes that TPLA's design (full-head access to partitioned latent) is structurally necessary for preserving accuracy during conversion — GLA's restricted per-head latent access causes catastrophic failure. The prefill–decode separation is convincingly shown to be more effective than short-text fine-tuning for preserving long-context performance.

The paper's weaker points are: (1) the throughput measurements exclude MoE routing, overstating the end-to-end practical benefit; (2) the absence of comparisons against KV cache quantization and token eviction methods leaves the relative advantage unclear; (3) the lack of domain diversity in evaluation (no code, math, or instruction-following benchmarks) limits confidence in the method's universality; (4) the PCA calibration cost and its sensitivity to calibration data choice are not characterized; (5) the $g>2$ limitation is acknowledged but not explored, bounding the method's applicability to higher TP degrees. The paper's strongest result — that training-free PD-sep. nearly matches original MLA accuracy across diverse tasks — is robustly supported by the reported numbers and represents a genuinely practical advance for deploying MLA-based models under tensor parallelism.

6. Limitations and Trade-offs

Limitation 1: PCA Reparameterization Breaks Down When Partitioning Into More Than Two Groups

The assumption or constraint. TPLA's reparameterization strategy — the mechanism that makes training-free conversion possible — depends on an orthogonal transform that concentrates activation variance so that the first half of dimensions captures nearly all information, making the second half a negligible perturbation. This works for $g=2$ groups (two GPUs), but the paper acknowledges in Section 6 that PCA-based reparameterization cannot scale to larger numbers of latent partitions:

"PCA concentrates most of the data's informative content in the first few dimensions, which provide a representative summary of the global structure. In contrast, the later dimensions primarily capture negligible noise and minor variations that contribute minimally to the overall representation. Consequently, TPLA with group-partitions $g=2$ can achieve good performance, but when $g>2$, it probably fails to maintain effectiveness."

When the latent dimension must be split into three or more groups, the middle partitions contain moderate-variance dimensions. These are neither dominant enough to approximate the global computation with a simple scaling factor (as the first partition can), nor negligible enough to treat as zero-mean noise (as the last partition can be in the $g=2$ case). The conditions that make per-shard RMSNorm and softmax approximations work — proportionality constants $\alpha$ and $\beta$ derived from eigenvalue ratios — have no clean generalization to an arbitrary number of partitions where each shard captures an intermediate fraction of variance.

The consequence. TPLA cannot effectively reduce per-device KV cache beyond a factor of approximately 2× regardless of the tensor parallelism degree. For a deployment requiring TP=4 or TP=8 — which is realistic for models like DeepSeek-V3 (685B parameters) or Kimi-K2 (1T parameters) at high batch sizes or long contexts — the KV cache per device bottoms out at $2d_h + d_r$ dimensions (320 for DeepSeek-V3). The memory bandwidth bottleneck is alleviated only once, when going from 1 to 2 devices; adding more devices provides no additional cache reduction. The paper describes a hybrid approach for TP>2 that splits heads within each latent group (Section 5.4.1: "we further split heads in addition to halving the latent dimension; for example, with TP=4 on Kimi-K2-TPLA, we use 32 heads × 2 per device with a 320-dimensional latent per head"), but this keeps the per-device cache size constant — the additional parallelism reduces compute, not memory traffic. Since decoding throughput under memory-bound conditions is determined by cache size, not head count, the speedup ceiling is approximately 2× regardless of how many GPUs are available.

What evidence exists in the paper. The paper provides no experiments with $g>2$. All throughput measurements (Figure 3) use two GPUs with the latent dimension split into exactly two groups. The speedups of 1.79× (DeepSeek-V3) and 1.93× (Kimi-K2) at 32K context length represent the maximum achievable benefit under the current reparameterization scheme. The paper does not report accuracy or throughput for $g=3$, $g=4$, or any other multi-group configuration that would stress-test the PCA approach's limits. The discussion of $g>2$ in Section 4.4 is purely algebraic (describing the general sharding pattern) without experimental validation.

Mitigation status. The paper suggests that Hadamard-based transforms — which balance information uniformly rather than concentrating it — may be more suitable for $g>2$ partitioning, since each shard would capture an equal fraction of the norm and the proportionality constants would remain balanced regardless of the number of groups. However, the paper also shows that Hadamard transforms fail to satisfy the softmax slicing condition (Condition 2) due to mixed-sign contributions that cancel globally but not locally, and acknowledges this as unresolved: "In future work, we will design and evaluate optimized Hadamard-like orthogonal matrices to balance softmax slicing, thereby improving both robustness and scalability" (Section 6). Until this is solved, TPLA's practical deployment is constrained to TP=2 configurations for the latent attention component. This is partially mitigated by the fact that head-level TP can provide additional parallelism within each latent group, but the memory bandwidth benefit — which is the primary motivation for TPLA — does not scale beyond two devices.

---

Limitation 2: Throughput Speedups Are Measured With MoE Routing Removed, Not on the Full Production Models

The assumption or constraint. The headline throughput improvements — 1.79× for DeepSeek-V3-0324 and 1.93× for Kimi-K2-Base at 32K context length — are measured on versions of these models with the Mixture-of-Experts (MoE) routing removed. The paper states this explicitly in Section 5.4.1:

"Because these models are extremely large and Mixture-of-Experts (MoE) routing can confound attention-speed effects, we remove MoE for timing."

MoE routing is a defining architectural feature of both models — DeepSeek-V3 uses a mixture of 256 experts with 8 experts activated per token, and the routing and expert computation constitute a substantial fraction of total inference time. The attention component that TPLA accelerates is only one part of the full model; the MoE layers, MLP layers, embedding lookups, and layer norms are unaffected by the TPLA reformulation.

The consequence. The reported speedups of 1.79×–1.93× apply to the attention component only, not to end-to-end inference throughput. A practitioner deploying the full DeepSeek-V3 model with MoE enabled will see a smaller overall speedup, attenuated by the fraction of total inference time spent in attention. The magnitude of this attenuation depends on context length (attention cost grows linearly with sequence length in decoding, non-linearly in prefill, while MoE cost is roughly constant per token) and on the specific MoE configuration. At short context lengths where attention is a minority of total FLOPs, the end-to-end benefit could be quite modest — perhaps 1.2×–1.4× rather than 1.8×. The paper provides no end-to-end measurements on the full models, so the practitioner cannot determine the actual deployment benefit without running their own benchmarks.

Additionally, TPLA introduces an all-reduce operation on the attention outputs (Equation 7) that was not present in standard MLA-TP (which also uses all-reduce for head outputs, but the communication patterns differ). In a full MoE deployment, there are already all-to-all communications for expert routing; TPLA's additional collective operations compete for the same interconnect bandwidth. The paper does not analyze how TPLA's communication overhead interacts with MoE communication in a full model, which could further erode the net throughput gain beyond what a simple FLOPs-count attenuation would predict.

What evidence exists in the paper. The paper provides thorough attention-component measurements in Figure 3 (throughput) and Figure 4 (TTFT), both with MoE removed. The complexity analysis in Section 4.5 compares attention-specific FLOPs and memory traffic between MLA-TP and TPLA-TP. No end-to-end measurements on the full MoE models are reported. The paper does not discuss how to estimate end-to-end speedup from the attention-component numbers or provide the ratio of attention time to total inference time for the evaluated models, which would allow a practitioner to approximate the real-world benefit.

Mitigation status. The paper is explicit about the MoE removal, so this is a methodological choice rather than an oversight. The authors clearly state their intent is to "isolate attention-speed effects." However, the abstract and introduction present the speedup numbers without qualification (e.g., "we achieve 1.79× and 1.93× speedups"), which could mislead a reader who does not carefully check the experimental setup. The practical mitigation — running end-to-end benchmarks on the full models — is left to the practitioner. A partial mitigation is theoretical: since decoding is memory-bandwidth-bound, and the KV cache is the dominant memory traffic source at long contexts, the attention speedup should be the dominant factor in end-to-end speedup at sufficiently long sequences. But the threshold at which this dominance kicks in is model-specific and not characterized.

---

Limitation 3: No Comparison Against Quantization or Token Eviction Methods That Achieve Similar Memory Reduction

The assumption or constraint. TPLA reduces the per-device KV cache by restructuring the attention computation itself — partitioning the latent dimension across GPUs, absorbing orthogonal transforms into weight matrices, and accepting approximation error in softmax and RMSNorm. The paper frames this as the solution to MLA's tensor parallelism problem and compares it against two baselines: standard MLA with replicated cache (the status quo problem) and Grouped Latent Attention (the prior architectural proposal). It does not compare against a much simpler and widely deployed class of solutions: post-hoc KV cache compression through quantization (int8, int4), token eviction (H2O, SnapKV), or token merging.

The paper acknowledges these methods exist in its related work (Section 2):

"Although effective, these approaches inevitably discard or alter information in the KV cache and can degrade model performance. In contrast, TPLA leaves the KV contents intact: it reduces the amount of cache each device must hold so the model retains full information while alleviating memory pressure."

The consequence. This framing implies that TPLA is superior because it preserves information, while post-hoc methods degrade accuracy. But the paper's own evidence shows that TPLA also degrades accuracy — the PD-sep. variant loses 2.15% on LongBench for DeepSeek-V3 (Table 2). The relevant question for a practitioner is: at matched accuracy degradation, which approach yields better throughput? If int8 KV cache quantization applied to standard MLA achieves 1.8× memory reduction with 2% accuracy loss and requires zero weight modification, zero calibration data, and zero code changes beyond the quantization kernel, the case for TPLA's complexity — weight reparameterization, PCA calibration, phase-dependent attention switching — weakens considerably. Worse, if quantization and TPLA are complementary (a 1.8× reduction from TPLA multiplied by a 2× reduction from int8 quantization yields 3.6× total), the paper provides no evidence that they compose without compounding accuracy loss.

Similarly, token eviction methods like SnapKV prune low-importance tokens from the cache, directly reducing cache size without changing the attention computation. For MLA models, these could be applied out-of-the-box with no structural changes. The paper does not test whether TPLA plus token eviction yields additive benefits or whether token eviction alone matches TPLA's speedup at lower implementation complexity.

What evidence exists in the paper. None. The paper provides no experiments comparing TPLA against quantized MLA, no experiments applying token eviction to either MLA or TPLA, and no analysis of whether TPLA's memory reduction is orthogonal to or redundant with post-hoc compression methods. The related work discusses these methods but the experimental section treats them as out of scope.

Mitigation status. Not addressed. The paper's claim that "TPLA leaves the KV contents intact" is technically true — the cache values are not truncated or quantized — but this misses the point that the attention computation is changed (softmax is decomposed, RMSNorm is approximated), which introduces its own form of information loss. A fair comparison would match accuracy degradation across methods and compare the resulting throughput, or would demonstrate that TPLA's approach composes multiplicatively with quantization (suggesting the two attack different bottlenecks). The absence of these comparisons means the practitioner cannot determine whether TPLA is the best available solution to the MLA-TP memory problem or simply a solution that works.

---

Limitation 4: Evaluation Restricted to Academic Benchmarks With No Instruction-Following, Code, Math, or Safety Assessment

The assumption or constraint. The paper evaluates TPLA exclusively on academic accuracy benchmarks: WikiText-2 perplexity, six commonsense reasoning datasets (MMLU, ARC, PIQA, HellaSwag, OpenBookQA, WinoGrande), and LongBench for long-context understanding. These benchmarks test factual knowledge, commonsense reasoning, and long-document comprehension — but they do not test the capabilities that make models like DeepSeek-V3 and Kimi-K2 useful in production: instruction-following, multi-turn conversation, code generation, mathematical reasoning, or tool use. The paper also provides no safety evaluation — no measurement of whether TPLA conversion affects refusal rates, toxicity, bias, or jailbreak susceptibility.

The core architectural change in TPLA — decomposing the softmax across two GPUs with per-shard normalization — introduces approximation error that is data-dependent. The quality of the approximation depends on how the attention scores distribute across the latent dimensions for a given input type. Inputs that produce sharply concentrated attention patterns (e.g., precise instruction-following requiring attention to specific constraint tokens) could be disproportionately affected by softmax slicing errors compared to inputs with diffuse attention patterns (e.g., general knowledge questions where many tokens are relevant).

The consequence. A practitioner deploying TPLA for a chat application, code assistant, or math tutor cannot rely on the paper's accuracy results to predict real-world performance. The 2.15% average drop on LongBench might translate to a 10% drop on code generation accuracy or a qualitative degradation in instruction adherence that frustrates users. More concerningly, safety alignment — which is typically applied post-training through RLHF or preference optimization and may depend on specific attention patterns for refusal triggers — could be silently degraded by the reparameterization and approximation changes, producing a model that is more likely to comply with harmful requests in ways that academic benchmarks cannot detect.

What evidence exists in the paper. LongBench provides the broadest coverage (21 tasks across 6 categories including QA, summarization, and few-shot learning), and the PD-sep. variant maintains strong performance across this diversity. This provides indirect evidence of robustness — if the approximation error were causing catastrophic failures on specific attention patterns, the LongBench average would likely show larger drops. But LongBench does not include code, math, multi-turn chat, or instruction-following formats, so the coverage is incomplete. No safety benchmarks are mentioned or evaluated.

Mitigation status. Not addressed. The paper does not acknowledge the gap between its evaluation suite and the models' intended use cases as a limitation. The choice of benchmarks is standard for academic ML papers but insufficient for a systems paper that proposes a drop-in replacement for production inference pipelines. The mitigation would be straightforward — evaluate on HumanEval or MBPP for code, GSM8K or MATH for reasoning, MT-Bench or AlpacaEval for instruction-following — but is left entirely to future work or to practitioners performing their own validation.

---

Limitation 5: The GLA Comparison Is Tested Only in Zero-Shot Conversion, Not From-Scratch Training, Weakening the Representational Capacity Claim

The assumption or constraint. One of the paper's central claims is that TPLA "preserves MLA's representational capacity" while GLA loses it because each attention head "only accesses half of the latent representation" (Section 1). The empirical support for this claim comes from a single experiment in Table 1: directly converting an MLA-pretrained DeepSeek-V2-Lite checkpoint to the GLA structural format with no retraining, which causes WikiText-2 perplexity to explode from 6.31 to 2212 — a catastrophic 350× degradation.

This experiment demonstrates that MLA-pretrained weights are incompatible with the GLA structure, but it does not test whether GLA as an architecture has lower representational capacity than TPLA. GLA was designed to be trained from scratch (Zadouri et al., 2025). The paper's claim is about inherent architectural capacity — "the reduction in KV cache size for single device comes at the cost of decreased representational capacity for each attention head" (Section 1) — but the evidence only tests how pre-trained MLA weights transfer to a GLA-shaped model, not the capacity of GLA when properly trained.

The consequence. A reader might conclude that GLA is fundamentally inferior to TPLA, and that future model training efforts should prefer TPLA over GLA. This conclusion is not supported by the evidence. The relevant comparison would be: train GLA and TPLA from scratch with matched compute budgets, then evaluate accuracy at matched per-device KV cache sizes. It is possible that GLA — which has fewer parameters per head (since each head accesses only half the latent dimension) and therefore lower computational cost — could achieve comparable accuracy to TPLA when both are trained in their native formats, because GLA's head groupings might learn to specialize in ways that compensate for the reduced per-head dimensionality. The paper provides no evidence either way.

The paper partially addresses this analytically in Section 4.4 by showing that TPLA can be formulated as GLA with doubled heads, meaning TPLA is structurally a superset of GLA — any computation GLA can perform, TPLA can also perform by setting the duplicated head weights appropriately. This is a valid structural argument that TPLA's representational capacity is at least as high as GLA's. But structural capacity is not the same as learnable capacity. The paper does not test whether the additional degrees of freedom in TPLA (the doubled heads) actually translate to better accuracy when trained from scratch, or whether they simply add optimization difficulty or overfitting risk that negates the theoretical advantage.

What evidence exists in the paper. Only the zero-shot MLA-to-GLA conversion in Table 1. The paper acknowledges in Section 6 that training TPLA from scratch is left to future work: "We will post-pretrain DeepSeek-V3, or train a TPLA-based model from scratch, to further demonstrate TPLA's excellent expressiveness." The GLA from-scratch comparison is not mentioned in this future work statement.

Mitigation status. Partially mitigated by the structural argument in Section 4.4, which provides a theoretical basis for TPLA's advantage independent of training regime. The paper is also transparent about not having trained TPLA from scratch. However, the strong language comparing representational capacity ("Unlike Grouped Latent Attention (GLA), every head in TPLA still leverages the full latent representation, maintaining stronger representational capacity" — Abstract) should be tempered by the acknowledgment that this claim is tested only in a conversion setting where GLA is set up to fail. A practitioner deciding between advocating for GLA or TPLA in a future model training project needs a from-scratch comparison that the paper does not provide.

---

Limitation 6: PCA Calibration Cost and Sensitivity Are Uncharacterized, Limiting Deployment Guidance

The assumption or constraint. TPLA's reparameterization requires computing PCA eigenvectors from a calibration dataset — the paper uses WikiText-2 throughout. This involves running the pretrained MLA model in forward-pass mode over the calibration corpus, collecting KV latent activations $F \in \mathbb{R}^{(B\cdot L) \times 4d_h}$ from every transformer layer, computing the covariance matrix and its eigenvalue decomposition for each layer independently. For a model like DeepSeek-V3 with 61 transformer layers, this means performing eigenvalue decomposition on 61 matrices of size $576 \times 576$ (the latent dimension), plus the forward passes to collect the activations.

The paper characterizes PD-sep. TPLA as "training-free" — which is accurate in the gradient sense — but the PCA calibration step is a non-trivial computational procedure with its own cost. The paper reports the alignment cost (100M tokens of SmolLM-Corpus for TPLA align) but not the PCA calibration cost (how many tokens, how many GPU-hours, for which model sizes). For a practitioner deploying TPLA, this missing cost matters: if PCA calibration requires 200 GPU-hours on the target model and must be redone after any fine-tuning (since the activation statistics change), the "training-free" label is somewhat misleading — the conversion is not zero-cost.

Additionally, the paper provides no sensitivity analysis studying how the choice of calibration dataset affects conversion quality. The PCA eigenvectors are specific to the activation statistics of the calibration corpus. If a practitioner's deployment data has different statistical properties — different languages, different domains (code, medicine, law), different input lengths — the variance fractions $\alpha$ and $\beta$ derived from WikiText-2 may not accurately represent the deployment distribution. This could cause the per-shard RMSNorm and softmax approximations to be less accurate in deployment than in evaluation, widening the accuracy gap beyond what the paper reports. The paper's LongBench results (which include Chinese and diverse task types) provide indirect evidence of reasonable generalization from WikiText-2 English text, but a direct ablation comparing calibration datasets would quantify how much the results depend on this choice.

What evidence exists in the paper. The paper uses WikiText-2 for PCA calibration and SmolLM-Corpus for alignment. No calibration dataset ablation is reported. No sensitivity analysis measures how conversion accuracy changes if PCA is computed from a different corpus or from a corpus with different length distributions. The calibration computational cost is not reported. The paper does not discuss whether the PCA transform needs recomputation after model fine-tuning, though this follows logically from the dependence on activation statistics.

Mitigation status. The paper does not address this as a limitation. The strong LongBench results (54.44% for TPLA PD-sep. vs. 56.59% for MLA on DeepSeek-V3) suggest the WikiText-2-derived PCA generalizes reasonably across task types, but the paper provides no engineering guidance on calibration data selection. A practitioner is left to assume that WikiText-2 is sufficient, or to run their own sensitivity experiments. Given that the calibration cost scales with model size (requiring full forward passes of the target model), testing multiple calibration datasets could itself be expensive, creating a chicken-and-egg problem: you need to calibrate to evaluate conversion quality, but evaluating conversion quality with different calibrations is expensive. The paper's omission of this practical consideration makes deployment planning more uncertain than the headline accuracy numbers suggest.

7. Implications and Future Directions

How This Work Changes the Landscape

TPLA is best understood not as a paradigm shift, but as a precision diagnostic and structural solution that resolves a previously unarticulated tension between two widely deployed inference optimizations. The paper does not propose a fundamentally new attention mechanism — it inherits MLA's architecture and GLA's sharding pattern — but it makes a contribution of a different kind: it identifies exactly why MLA and tensor parallelism clash, and provides an engineered pathway to make them compatible without retraining. This is systems research at its most useful: the value is not in novelty of concept but in the specificity of the problem diagnosis and the practicality of the fix.

The landscape change occurs along three axes. First, it recharacterizes MLA's efficiency claim with a sharp boundary condition. Prior to this work, MLA was understood as a KV cache compression technique — end of story. The paper shows that this understanding is incomplete: MLA's efficiency must be evaluated under the parallelism strategy used at deployment. The cache goes from 576 dimensions (MLA's celebrated efficiency) to effectively larger per device than GQA under sufficient TP. This reframing makes parallelism-aware cache compression a first-class design criterion for future attention mechanisms. Future MLA variants or MLA-like low-rank attention schemes should be designed from the start with a sharding plan that distributes the compressed representation across devices, not just with the goal of minimizing total cache size on a single GPU.

Second, it demonstrates that weight-space reparameterization can substitute for training when adapting model architectures to new inference constraints. The standard response to an architecture change that introduces approximation error is fine-tuning — retrain the model a bit to adapt to the new computation graph. TPLA's PCA-based reparameterization plus prefill–decode separation achieves near-original accuracy with zero gradient updates by instead pre-rotating the weight space so that the subsequent decomposition introduces minimal error. This reparameterization-as-error-control pattern is not widely exploited in the inference optimization literature, which tends to favor post-hoc quantization or pruning with optional fine-tuning recovery. TPLA's success — 2.15% average degradation on LongBench for DeepSeek-V3 without training (Table 2) — suggests that many inference-time decompositions that appear lossy might become near-lossless if approached as a basis-change problem rather than an approximation-then-recover problem. This opens a methodological door: for any non-linearity that must be decomposed across devices (softmax, normalization, activation functions), ask first whether an orthogonal transform absorbed into adjacent weights can rebalance the computation so that per-shard operations better approximate the global operation.

Third, it reconciles the apparent contradiction between GLA's proposal and the realities of existing model ecosystems. GLA (Zadouri et al., 2025) correctly diagnosed MLA's TP limitation and proposed latent-dimension partitioning as the solution. But GLA's prescription — train a new model from scratch — is incompatible with the practical reality that the most capable open-weight models using MLA (DeepSeek-V3, Kimi-K2) have already been trained at enormous cost. TPLA shows that GLA's core insight (partition the latent dimension) can be realized on existing checkpoints through reparameterization, without GLA's restrictive head-to-latent grouping that causes the 350× perplexity explosion in Table 1. This shifts the conversation from "we need new architectures for TP-efficient inference" to "we can adapt existing architectures for TP-efficient inference." Given the enormous inertia behind pretrained model ecosystems, this adaptability argument is pragmatically compelling.

The research directions that become more attractive include: basis-change reparameterization for other attention variants (e.g., applying orthogonal transforms to sliding window or linear attention mechanisms to improve their TP sharding), combining TPLA's structural memory reduction with quantization (attacking memory from complementary angles), and designing attention mechanisms that are natively "TPLA-ready" — trained with the latent dimension already partitioned across devices. The research directions that become less attractive include: proposals for entirely new attention architectures that require from-scratch training without a clear migration path from existing MLA models (since TPLA shows adaptation is often possible), and head-level TP optimizations for MLA (since the paper shows these don't reduce cache memory, the primary bottleneck).

Follow-Up Research This Work Enables

1. Train TPLA from scratch and compare against GLA at matched compute budgets. The paper demonstrates that TPLA can convert MLA checkpoints with minimal loss, but does not test whether TPLA as a training architecture outperforms GLA. A convincing follow-up would pretrain two model variants — TPLA (with $g=2$ latent groups and full head replication on each device) and GLA (with $g=2$ latent groups and restricted head access per group) — at the same parameter count and training token budget, then evaluate on standard benchmarks. The hypothesis is that TPLA's full-head access to the partitioned latent dimension yields better accuracy at matched KV cache size, which the structural argument in Section 4.4 predicts but does not empirically validate. A negative result (GLA matches TPLA when both are properly trained, because GLA's head specialization compensates for reduced per-head dimensionality) would be equally informative, suggesting that the conversion advantage — not an inherent capacity advantage — is TPLA's primary contribution.

2. Design and benchmark optimized Hadamard-like transforms for softmax slicing with $g>2$ groups. The paper identifies that PCA concentrates variance in the first few dimensions, making it unsuitable for partitioning the latent space into three or more roughly equal-capacity groups (Section 6). Hadamard transforms balance norms perfectly but fail the softmax slicing condition due to signed contributions that cancel globally but not locally (counterexample in Section 4.3.1). A targeted investigation would search over a parameterized family of orthogonal matrices — perhaps block-diagonal Hadamard variants with sign-constrained blocks, or learned orthogonal transforms optimized via gradient descent on a softmax approximation loss — to find transforms that simultaneously satisfy both RMSNorm balancing (uniform per-shard norm) and softmax balancing (per-shard attention scores proportional to global scores up to multiplicative constants). The evaluation would measure accuracy at $g=2, 4, 8$ on a fixed benchmark (e.g., MMLU + LongBench) using DeepSeek-V2-Lite, comparing the optimized transform against PCA, vanilla Hadamard, and random orthogonal baselines. Success would directly extend TPLA's memory bandwidth benefits to higher TP degrees, where the per-device cache could be reduced by factors of 4× or 8× rather than being capped at ~2×.

3. Combine TPLA with KV cache quantization and measure composability. The paper avoids quantization entirely, and its related work frames quantization as an information-discarding alternative rather than a complementary technique. A practical follow-up would apply standard int4 or int8 KV cache quantization (e.g., KIVI, KVQuant) on top of TPLA's already-reduced per-device cache, measuring end-to-end throughput on the full MoE models (DeepSeek-V3 with routing enabled) and accuracy degradation at each combination. The key question: does the 1.8× memory reduction from TPLA multiply with the 2× reduction from int8 quantization to yield ~3.6× total cache compression, or do the approximation errors compound non-linearly such that TPLA + int4 is worse than either alone at matched accuracy? If the benefits are multiplicative, the case for TPLA strengthens enormously — a practitioner gets TPLA's structural advantage plus quantization's compression advantage. If they interfere (e.g., TPLA's softmax slicing error is amplified when the already-approximate attention scores operate over quantized cache values), TPLA and quantization become alternatives rather than complements, and practitioners must choose based on accuracy-throughput tradeoff curves that this follow-up would map.

4. Characterize sensitivity to calibration data distribution with domain shift experiments. TPLA's PCA reparameterization depends on activation statistics from a calibration corpus (WikiText-2 in the paper). A rigorous follow-up would ablate this choice systematically: compute PCA eigenvectors from several corpora with distinct statistical properties (e.g., WikiText-2 for general English, The Stack for code, a multilingual corpus like mC4, and a math corpus like OpenWebMath), convert DeepSeek-V3 using each, and evaluate on both in-domain and out-of-domain benchmarks. Specifically, evaluate code-generation accuracy (HumanEval, MBPP) after calibration on WikiText-2 vs. The Stack; evaluate multilingual reasoning after calibration on WikiText-2 vs. mC4; evaluate mathematical reasoning (GSM8K) after calibration on WikiText-2 vs. OpenWebMath. The output is a domain-specific accuracy matrix that tells practitioners whether they need to match their calibration data to their deployment domain, or whether the WikiText-2 PCA generalizes universally. A negative result (large accuracy swings depending on calibration data) would imply that TPLA requires careful calibration data selection — a practical burden not acknowledged in the paper. A positive result (minimal sensitivity) strengthens TPLA's "training-free" claim by showing that the calibration step is robust to data choice.

5. Evaluate TPLA on instruction-following, safety, and structured reasoning tasks. The paper's evaluation suite covers language modeling and commonsense reasoning, completely omitting the task types that define production LLM usage: multi-turn conversation, instruction adherence, code generation, mathematical reasoning, and safety alignment. A deployment-motivated follow-up would run TPLA-converted DeepSeek-V3 (PD-sep.) on MT-Bench or AlpacaEval for instruction-following quality, HumanEval and MBPP for code, GSM8K and MATH for reasoning, and standard safety benchmarks (e.g., TruthfulQA, ToxiGen, refusal rate on harmful prompts) to detect any degradation in alignment. The concern is that softmax slicing error — which affects how attention is distributed across tokens — might disproportionately impact tasks requiring precise token-level attention, such as following multi-constraint instructions or attending to specific variable names in code. A positive result (no significant degradation on any of these) would substantially strengthen TPLA's claim to being a drop-in replacement for production MLA inference. A negative result (degradation on, say, code generation despite strong LongBench performance) would characterize TPLA's domain of applicability more precisely and guide practitioners on which workloads benefit from TPLA and which should retain standard MLA or use alternative optimizations.

6. Investigate whether post-training can close the remaining 2.15% LongBench gap more effectively than short-text alignment. Table 2 shows that TPLA (align), which uses 100M tokens of mostly short concatenated texts from SmolLM-Corpus, underperforms the training-free PD-sep. variant on DeepSeek-V2-Lite LongBench (38.53% vs. 40.02%) and only partially recovers on DeepSeek-V3 (53.48% vs. 54.44% for PD-sep.). This suggests that the residual approximation error in TPLA is concentrated in long-context attention patterns that short-text fine-tuning does not address. A targeted follow-up would fine-tune TPLA-converted models on a long-context corpus (e.g., sequences of 32K–128K tokens from Books3, long-form QA datasets, or synthetic long-context data) and measure whether the gap to original MLA can be closed to within, say, 0.5%. This would test whether the remaining degradation is a fundamental consequence of the softmax decomposition (unrecoverable without changing the attention structure) or a trainable adaptation that the paper's alignment recipe simply failed to trigger due to data mismatch. The result directly informs whether TPLA can realistically match MLA accuracy in all deployment scenarios with modest additional compute, or whether a fundamental accuracy-throughput tradeoff remains inherent to the decomposition.

Practical Applications and Downstream Use Cases

1. High-throughput API serving for MLA-based models under tensor parallelism. The most immediate use case is inference serving for DeepSeek-V3, Kimi-K2, and future MLA-pretrained models when deployed across multiple GPUs. At 32K context length with TP=2, TPLA's 1.79× (DeepSeek-V3) and 1.93× (Kimi-K2) decoding throughput improvements (Figure 3) translate directly to serving more requests per second without increasing GPU count. For an inference provider operating at scale — where each 1% throughput improvement translates to measurable cost savings — the training-free PD-sep. variant enables immediate deployment on existing checkpoints with only a 2.15% average accuracy cost on long-context tasks (Table 2). The deployment workflow is: (1) run PCA calibration on the target model using a representative text corpus (~hours on ~8 GPUs for a 685B model), (2) absorb orthogonal transforms into the weight matrices via the reparameterization equations (Equations 17, 21), (3) configure the inference server to use unsliced MLA during prefill and TPLA with latent-dimension partitioning during decoding. The key practical consideration is that the MoE routing overhead (not accelerated by TPLA) will dilute the end-to-end speedup — the attention-component 1.8× translates to a smaller but still meaningful end-to-end gain whose exact magnitude depends on the attention-to-MoE compute ratio at the target context length.

2. Long-context batch inference on GPU-limited hardware. For research labs or smaller organizations that need to run long-context inference on DeepSeek-class models but are GPU-constrained, TPLA enables fitting larger effective batch sizes or longer contexts within the same GPU memory budget. The per-device KV cache drops from 576 to 320 dimensions per token for DeepSeek-V3, a 1.8× reduction. At 32K context length, this is the difference between the KV cache fitting in HBM with room for large batch sizes versus hitting out-of-memory limits that force smaller batches or shorter contexts. The practical benefit is not just throughput (more tokens per second) but feasibility — workloads that exceed GPU memory without TPLA become runnable with TPLA, without model downgrades or aggressive KV cache eviction that could harm accuracy. The training-free PD-sep. variant is particularly valuable here because the deploying team may not have the resources for fine-tuning, and the benchmark coverage (commonsense + long-context) provides sufficient confidence for research use cases where minor accuracy degradation is acceptable.

3. Self-improvement and synthetic data generation pipelines using MLA models. When using DeepSeek-V3 or Kimi-K2 to generate training data (e.g., for distillation, instruction-tuning dataset creation, or iterative self-improvement loops), inference cost is the dominant expense. These pipelines often run thousands or millions of forward passes, and decoding throughput — not prefill latency — is the bottleneck. TPLA's throughput improvement applies directly: generating 1M tokens of synthetic data at 32K context length costs approximately 1.8× fewer GPU-hours with TPLA than with standard MLA under TP=2. The 2.15% LongBench accuracy drop (Table 2) must be evaluated against the task-specific accuracy requirements of the data-generation pipeline — for many use cases (creative writing, summarization, translation, general instruction generation), this level of degradation is likely acceptable in exchange for nearly halving the inference cost. A caveat: the paper does not evaluate instruction-following quality after conversion (Limitation 4), so a pipeline generating structured instruction-response pairs should validate output quality with a small pilot run before committing to TPLA for full-scale generation.

4. On-premise deployment of large MLA models for privacy-sensitive applications. Organizations that must run LLM inference on-premise (healthcare, finance, government) often face GPU memory constraints that limit the maximum usable model size or context length. TPLA's per-device KV cache reduction means that a fixed GPU configuration can support either (a) the same model with ~1.8× longer usable context windows, or (b) a larger batch size at the same context length, improving hardware utilization. The 1.4× prefill latency improvement from PD-sep. (Figure 4) additionally reduces Time to First Token, which is important for interactive applications where users wait for the model to process their input before generation begins. The training-free nature of PD-sep. TPLA is particularly valuable in regulated settings where any model modification (even fine-tuning) may require re-validation or re-certification — since PD-sep. introduces no gradient updates and preserves accuracy within ~1 percentage point on commonsense tasks (Table 1), it may fall into a different regulatory category than fine-tuned model variants, simplifying deployment approval processes.