TransMLA: Migrating GQA Models to MLA with Full DeepSeek Compatibility and Speedup

Thank you very much for your encouraging review and constructive feedback. Below we provide point-by-point responses to your comments. All suggested improvements will be incorporated into the camera-ready version.

Q1. While intuitive and empirically validated, the balancing method lacks theoretical guarantees or principled formulation.

A1. Balanced Key-Value PCA (BKV) is indeed a simple yet effective technique. While many prior works (Palu, MHA2MLA, X-EcoMLA) have not adopted this method, we highlight it as a novel contribution. Behind its apparent simplicity, BKV has a solid theoretical foundation. We prove that under a joint loss function minimizing the relative approximation errors of both $W_K$ and $W_V$ components: $\mathcal{L} = \|| \frac{W_K - W_K^{top-r}}{\mathrm{RMS}(W_K)} \||^2 +\|| \frac{W_V - W_V^{top-r}}{\mathrm{RMS}(W_V)} \||^2$, scaling $W_K$ and $W_V$ to the same magnitude: $W_{\text{scaled}}^{KV} =$W_K / \mathrm{RMS}(W_K), W_V / \mathrm{RMS}(W_V)$$ achieves the optimal low-rank approximation under this specific loss function. In our proof, we use $a$ and $b$ as scaling factors. We first show that the loss function depends only on the ratio $a:b$, not on the absolute values of $a$ and $b$. By setting $b=1$, the optimization simplifies to finding the optimal value for $a$. Next, we compute the derivative of the loss function with respect to $a$, and show that: $\frac{d\mathcal{L}(a, 1)}{da} = 0$ when $a = \mathrm{RMS}(W_K) / \mathrm{RMS}(W_V)$, through a series of substitutions. The main substitutions used are: $\||W\||^2 = \mathrm{tr}(W^T W) = \sum_{i=1}^{\mathrm{rank}(W)} \sigma_i^2, \quad \text{and} \quad \frac{d\sigma_i}{da} = u_i^T \frac{dW}{da} v_i$, where $u_i$, $\sigma_i$, and $v_i$ are the left singular vectors, singular values, and right singular vectors of matrix $W$, respectively. This proof provides a theoretical foundation and a more principled mathematical formulation for our Key-Value balancing strategy. The detailed proof and mathematical derivation will put in the appendix of the camera-ready version.

Q2. On the need for a calibration dataset and concerns in privacy-sensitive or black-box settings:

A2: Regarding privacy-sensitive datasets, please see our response to Reviewer qHce A2: we address this by using distilled data generated from the original model, which avoids direct access to private corpora.

As for black-box models, TransMLA is primarily intended to accelerate inference, which generally requires access to model internals. However, we will fully open-source TransMLA, including conversion, training, and inference code, as well as converted checkpoints. This allows model providers to apply the method internally without data sharing. Alternatively, they may use our converted foundation models and fine-tune them on in-house data.

In addition, we also provide alternative implementations of RoRoPE, FreqFold, and BKV that do not require calibration data. Instead of performing PCA on activations, these versions apply SVD directly to the weight matrices. This option will be included in our open-sourced codebase for users who prefer not to rely on calibration datasets.

Q3. On clarity in exposition, especially for RoPE slicing and rotation math in Section 4.2:

A3. Thank you for this valuable suggestion. Although we made a conscious effort to ensure readability during writing, we recognize that several theoretical components in TransMLA—especially RoRoPE—remain abstract and challenging to follow. We have since redesigned the conversion pipeline diagram to include each component’s purpose. Additionally, we created intuitive frequency-domain visualizations to explain both RoRoPE and FreqFold. These updates significantly improve clarity and allow readers to grasp the motivation even without diving into all the formulas. These diagrams will be included in the camera-ready version.

Q4. On releasing converted MLA models such as LLaMA2-7B:

A4. Absolutely. We will release all code for conversion, training, and inference, as well as checkpoints for converted models. In the future, we plan to support more models and more tasks, providing the community with faster and stronger pretrained models.

Q5. On the sensitivity of Balanced KV to the scaling constant and potential for optimization:

A5. To investigate this, we conducted an ablation study as shown in Table 8. On LLaMA-3-8B, completely removing Balanced KV results in a 3.8-point increase in perplexity, highlighting its critical role. We also observed that model performance is sensitive to the scaling coefficient used in Balanced KV. The default value in our paper already yields strong overall results. However, following your suggestion, we performed a per-layer grid search over the range [0.3, 2.0] with 20 intervals, aiming to minimize the output difference from the original model. This tuning led to a 1.2-point reduction in perplexity compared to the default, suggesting that layer-wise coefficient optimization is both feasible and beneficial for further improving performance.

Table 8. Ablation study on the impact of the Balance KV scaling factor on the performance of the converted model. In the main paper, the default factor is computed as the average L2 norm of $K_{\text{RoPE}}$ divided by the average L2 norm of V.

Thank you for the insightful questions. Below we provide point-by-point responses to your comments. All suggested improvements will be incorporated into the camera-ready version.

Q1. The experiments focus on MHA models, but most current models use GQA, which is harder to compress.

A1. As stated in lines 29–36 of the original paper, both MHA and MQA can be seen as special cases of GQA. Specifically, when the number of query heads equals the number of key-value heads in GQA, it effectively becomes an MHA. Our goal is to compare TransMLA with prior relevant work under consistent and fair experimental settings. To this end, we selected the recently published MHA2MLA [1] as our baseline, as it also targets RoPE decoupling and KV cache compression. Following their experimental setup, we used two MHA-based models—SmolLM-1.7B and LLaMA-2-7B. Moreover, TransMLA is fully compatible with GQA, all of our ablation studies were performed on LLaMA-3-8B, a GQA-based model.

Nevertheless, as you suggested, we extended our evaluation to several mainstream GQA models. The results are reported in Table 1.

Table 1. Evaluated with Wikitext-2 Perplexity (PPL) of Converting GQA Models to MLA (lower is better). Original PPL refers to the perplexity of the original GQA model without conversion. Decoupled RoPE PPL represents the result after applying RoRoPE and FreqFold to reduce decoupled RoPE loss. Compressed KV PPL indicates the perplexity when the KV cache is compressed to 576 with the help of Balance-KV.

From Table 1, we can observe that for GQA models—ranging from 7B to 70B, whether dense models or MoE models—TransMLA consistently incurs only minimal performance loss. It’s worth noting that MLA is specifically optimized for a latent dimension of 576, which we adopt as a unified compression target across all models. Since GQA models typically maintain smaller KV caches than MHA, converting them to MLA results in less information loss, making the transition both effective and efficient. Due to time constraints, we only present the results of the conversion without additional training. We plan to include the full alignment performance for all models in the camera-ready version.

Q2. The restoration experiment for LLaMA-2-7B uses different data from its original training. Could the improvements come from additional learning rather than true recovery?

A2. As mentioned in A1, our work follows the experimental setup of MHA2MLA and uses the smollm-corpus for pretraining. We noted that the pretraining corpus for LLaMA-2-7B has not been publicly released. To address this concern, we adopted the data distillation approach proposed in the LLM-QAT [2] paper, generating a new dataset from the original model to guide the distillation of the converted model. For each model, we distilled datasets containing 14,000 samples, with each sample consisting of 2,048 tokens. These datasets were used to distill the converted model by minimizing the error between the output of each converted MLA attention module and the original model's corresponding GQA outputs, as well as aligning with the final token prediction. Each attention module was trained for 10 epochs, and we perform hyperparameter tuning over 10 evenly spaced learning rates in the range [1e-3, 2e-2]. The results are presented in Table 2.

Table 2. Commonsense reasoning accuracy for distilling converted MLA-based models from original GQA models

According to Table 2, aligning the converted models using only the distilled data from the original models allows them to closely match the performance of the original models. This confirms that the observed performance improvements are not due to additional learning but reflect genuine restoration. We plan to distill more data and complete experiments on additional models and compression rates in the future.

[1] Towards Economical Inference: Enabling DeepSeek's Multi-Head Latent Attention in Any Transformer-based LLMs (ACL 2025)

[2] LLM-QAT: Data-Free Quantization Aware Training for Large Language Models (ACL Findings 2024)

I would like to reiterate that my concern is not with the fundamental incompatibility between GQA and TransMLA, but rather with the trade-off between compression rate and final performance. That said, the authors have provided key results on the main GQA-based model, which I believe sufficiently address the logical gap previously present in the manuscript. I will revise my score accordingly, pending further consideration of the appropriate adjustment. I wish the authors all the best.

We sincerely thank you for your detailed questions and constructive feedback. Below, we address each of your concerns and will incorporate the necessary updates in the camera-ready version.

Q1. Calibration dataset for PCA: effect of size and dataset choice

A1. As per your suggestion, we conducted an ablation study on the calibration dataset. As shown in Table 4, the number of calibration samples does affect the conversion quality, with 64 samples yielding the best performance. Using different datasets for calibration also impacts results slightly, but all perform better than using randomly generated tokens.

Table 4. Ablation study of using different calibration data sources and dataset sizes for converting LLaMA-3-8B to MLA.

Q2. Lack of ablation studies (e.g., commonsense reasoning performance for RoRoPE without FreqFold)

A2. We have presented ablations on perplexity (PPL) in Figure 2 and Figure 3 of our submission for RoRoPE, FreqFold, and BalanceKV. Following your suggestion, we additionally evaluated commonsense reasoning accuracy, and the results are shown in Table 5:

• Without RoRoPE, accuracy drops by ~20% compared to the best configuration.
• Without FreqFold (e.g. FreqFold = 1), accuracy drops by ~6.5%.
• Without BalanceKV, accuracy drops by ~10%.

These results highlight the importance of each component in maintaining strong downstream performance.

Table 5. Ablation study of removing each method and its impact on commonsense reasoning performance of the converted model without training.

Q3. Experiments on Qwen / Mistral

A3. As mentioned in Reviewer qHce A1, our work follows the experimental setup of MHA2MLA and uses the smollm-corpus for pretraining SmolLM-1.7B and LLaMA-2-7B. To answer your question, we converted Qwen and Mixtral to MLA and conducted distillation experiments to recover the performance. The results are shown in the following table.

Q4. Comparison with KVQuant and Palu

A4. KVQuant is discussed in Line 34 of our original paper, and Palu is discussed in Lines 109–114. As further explained in Reviewer i2PF’s comments A1 (second paragraph) and A3.2, TransMLA is orthogonal to KVQuant and can be combined with it. Palu avoids decoupling RoPE by compressing only modules that do not include positional encoding. In their Figure 5, they show that without RoPE, absorb enables a 6.17× speedup comparing to MHA module at a 64K context length. However, when RoPE is present, absorb is disabled, and the speedup drops to 2.91×. This highlights the significance of our contribution, which enables low-rank compression of RoPE—making absorb usable and thus enabling more efficient inference.

Q5. Test on LongBench benchmark

Following your suggestion, we select Qwen-2.5-7B (with max position embeddings of 131072) and convert it into an MLA model. The original model is then used to distill the converted one. We report the performance of both the original and the converted (distilled) models on multiple long-context tasks from the LongBench in Table 6.

Table 6. Performance on multiple long-text tasks and speed (tested on 16k-length context) of Qwen after converting GQA to MLA and applying further distillation, as described in Reviewer qHce A2.

Q6. Figure 1 is not well-explained.

A6. Thank you for pointing this out. We realized that Figure 1 contained too much information, so we have split it into three separate parts: an overall workflow diagram, a frequency-domain transformation illustration for RoRoPE, and a schematic of the FreqFold approximation. Each annotated component in the figures is now explained step-by-step in the Method section. These updates will be included in the camera-ready version.

Q7. How does proving that MLA has stronger expressive capacity than GQA given the same KV cache size or parameters help support the contributions of this paper?

A7: Please refer to Reviewer i2PF A1.1 (first paragraph) and A3.1. “MLA consistently exhibits greater expressive capacity than GQA given the same KV cache size” is a central motivation of this conversion.

Q8. MLA seems clearly stronger than GQA as there are more parameters.

A8. More parameters do not necessarily imply stronger expressive capacity across different model architectures. For example, as noted in the original DeepSeek paper, MHA with more parameters demonstrates lower expressive capacity than MLA. One of the key contributions of our paper is a constructive proof showing that, under equal KV cache size, all GQA can be converted into MLA, but the reverse is not true. This provides a formal guarantee that MLA always possesses stronger expressive capacity than GQA.

Q9. Any other study comparing the expressive capacity of GQA and MLA.

A9. Since the introduction of MLA last year, several studies have proposed hypotheses and conducted experiments to explore the source of MLA’s strong performance. For example, Jianlin Su (the author of RoPE) discussed on his blog that increasing the head dimensions in MLA may contribute to its improved expressiveness.

Q10. Comparison with X-EcoMLA

A10. X-EcoMLA removes all positional encodings in the low-rank (nope) part and averages across all heads to extract 64-dimensional RoPE embeddings. This approach lacks theoretical justification for its approximation validity. Both the rope and nope parts of their transformation introduce significant performance drop. As X-EcoMLA is not open-sourced yet, we were unable to reproduce and benchmark their method. Once their code is available, we will provide a direct comparison in the camera-ready version.

Q11. I am wondering how the authors benefit from the DeepSeek optimizations for the inference speed measurement. Is the FlashMLA adopted in this experiment? It would be beneficial to add more details for the inference optimization.

A11. FlashMLA supports only the Hopper architecture. In our original submission, we did not have access to Hopper GPUs, so we used VLLM with Triton backend for MLA, and FlashAttention backend for MHA. During the rebuttal period, we reran the experiments using Hopper GPUs with FlashMLA backend. As shown in Table 7, compared to the original GQA model on the same GPU, FlashMLA achieves up to an 11.32× speedup at 32K context length. The detailed information and the code will be publicly available on the camera ready paper.

Table 7. Throughput (tokens/s) comparison between LLaMA-2-7B and the converted MLA-based model, evaluated on Hopper architecture.

We first address your new questions one by one:

Q1. It is great to see that the proposed method could apply to Qwen. I am wondering how TransMLA approximates the qkv bias term in Qwen2.5-7B. Will that require significant changes to the algorithm?

A1. Thank you for your insightful question. TransMLA can be seamlessly applied to Qwen-2-7B. Regarding the bias in the qkv projection, we employ a simple yet effective approach: for any linear module with bias $Y = WX + b$, where $X \in \mathbf{R}^d$ is the input activation, we transform it into a bias-free form by defining $X' = (X, 1) \in \mathbf{R}^{d+1}$ and $W' = (W, b)$. This allows us to express the operation as $Y = W'X'$ without bias. Consequently, we can directly apply RoRoPE, FreqFold, and BKV techniques without additional complexity. We will provide more detailed explanations of this approach in the camera-ready version of our paper.

Q2. I appreciate that the authors did additional experiments with FlashMLA, and the results look good. One question I have is that FlashMLA seems to only support the DeepSeekV3 dimensions. When the number of heads and kv_rank of TransMLA are different from DeepSeekV3, how to support TransMLA with FlashMLA?

A2. Thank you for raising this important question. While FlashMLA was initially designed with DeepSeekV3's 128 heads in mind, it remains fully compatible with other head configurations— Kimi-K2 (64 heads) and Llama-2-7B-TransMLA (32 heads). Regarding kv_rank support, the current implementation does specialize for kv_rank=576. However, to enable flexibility while maintaining acceleration benefits, we've developed a grouped MLA approach: by partitioning the kv_cache into multiple groups (each with dimension 576), we can effectively support larger total kv_rank values. For example, a kv_rank of 1088 can be achieved through two 512-dimensional groups plus a 64-dimensional remainder carrying positional information. We will release the implementation details. This design ensures that TransMLA maintains both performance and adaptability across different model architectures.

Q3. This is a minor point. Regarding A8, I agree that more parameters do not translate to more expressive power, but it also depends on how many more parameters. When comparing GQA and MLA with the same KV cache budget, MLA seems to have more than 10% parameters, which is not negligible. I agree with the motivation of TransMLA, and I believe it is a useful technique. However, I think it is unfair and not very useful to analyze and compare the express power of GQA and MLA with the same KV cache size. When applying TransMLA to Llama or Qwen, the KV cache of TransMLA will always be smaller.

A3. Indeed, under the same KV cache budget, MLA does have more parameters than GQA, and in our proof, these additional parameters are a key reason for its improved capability. Since we aim to align inference speed when comparing expressive power, please refer to Reviewer i2PF's Table 3, where the relationship between KV cache and inference speed is more direct compared to the impact of parameter count.

Your new questions have been extremely helpful in improving our paper, and we believe they will make TransMLA more rigorous, readable, and reproducible. Have our responses fully addressed all your concerns? Based on the original text, the additional experiments and explanations provided during rebuttal, we would be very grateful if you consider raising your score.

Dear Reviewer KGKd,

Thank you for your professional question. Indeed, grouped TransMLA is not equivalent to the original TransMLA, as performing PCA in groups yields weaker results compared to joint PCA when retaining the same feature dimensions. We have experimentally verified this in the attached table: under the same KV cache conditions, the ungrouped approach outperforms the two-group version. However, the grouped method allows for retaining more KV cache, which is why it performs better than the original TransMLA with only a single 576-dimensional KV cache. Importantly, this grouped approach maintains compatibility with FlashMLA inference. We tested the speeds of MHA, GQA, two-group MLA, and single-group MLA (FlashMLA does not support inference with dimensions other than 576, so we did not compare the single-group 64+1024 MLA) under 16k context length. As shown in the table, the two-group MLA achieves a 6.2× speedup over MHA and a 1.58× speedup over GQA, while the single-group MLA is 1.9× faster than the two-group version.

We sincerely appreciate your satisfaction with our rebuttal and your consideration of raising the score. Given that the author-reviewer discussion will conclude in 12 hours, we would be immensely grateful if you could adjust your rating.

Respectfully,

Submission 3063 Authors

Q1. About the innovations proposed in this paper.

A1. This work makes solid theoretical and practical contributions:

First, we theoretically prove that MLA has a stronger expressive capacity than GQA under the same KV cache size (see Appendix A), addressing a key gap left by DeepSeek, which only empirically demonstrated MLA’s capabilities. This theoretical framework enables comparative analysis of the expressive power of major attention mechanisms, including MLA, MHA, GQA, and MQA and provides strong motivation for converting GQA to MLA.

Second, we are the first to prove sufficient conditions for performing invariant query/key transformations across RoPE (see Appendix B), namely, preserving the inner product between RoPE-applied query and key before and after transformation. This breakthrough is crucial for enabling MLA to benefit from inference acceleration while preserving position information, which was a significant challenge for prior works. For example, Palu [1] performs low-rank compression before applying RoPE, which not support absorption, requires recomputing the KV representations from the KV cache before computing attention. As a result, the acceleration benefit is limited. MHA2MLA [2] and X-EcoMLA [3] directly truncate RoPE in a non-equivalent manner, leading to significant loss of positional information and degraded performance.

Our theoretical result led to the development of RoRoPE, a constrained low-rank compression method, and FreqFold, a technique that enhances the expressive power of RoRoPE. Together, these innovations reduce the loss from decoupling RoPE to a minimal level, making GQA-to-MLA conversion with partial RoPE feasible.

Lastly, Balance-KV is a simple yet highly effective method for improving joint KV compression. To the best of our knowledge, this technique has not been adopted in prior related works [1,2,3], and we list it as a contribution. Please refer to Reviewer E7fv A1, where we provide the theoretical justification behind this method.

Q2.1. Discuss the performance differences, advantages and disadvantages of MHA and MLA

A2.1. In the Introduction, Related Work, and the section on Expressive Power Analysis, we have outlined the strengths and weaknesses of MHA, GQA, and MLA: 1) MHA: High expressive power, but slow inference and maximum KV cache. 2) GQA: Lower expressive power, but moderate inference speed. 3) MLA: A balance between expressive power and inference efficiency, with higher expressiveness than GQA under the same KV cache. MLA has already garnered significant attention in both model architecture design and AI infrastructure communities. Numerous improved variants such as GLA [4], MTLA [5], and GTA [6] have been proposed, and dedicated frameworks and hardware optimizations for MLA inference are actively under development. Given this momentum, our unified framework for analyzing the expressive capacity of multiple attention mechanisms becomes especially valuable and broadly applicable.

Q2.2. The innovations are overly dependent on Absorb operations

A2.2. The brilliance of MLA’s design, beyond its use of low-rank compression—a concept already explored in prior works—lies in the Absorb operation. This operation is central to MLA’s efficiency: during the training phase (which is compute-intensive), it allows attention computations to be performed using compact 192-dimensional representations per head, achieving high computational efficiency. During inference (which is memory-bound), it employs a shared 576-dimensional KV cache across heads by absorbing the key projections into the query projection matrices, substantially reducing memory access overhead on both high-bandwidth memory (HBM) and on-chip SRAM.

As demonstrated in Table 3 (rows 2–3), disabling the absorb operation causes MLA’s inference speed at a 16k context length to degrade by 16.7×, clearly showing its critical role. Consequently, this paper places strong emphasis on supporting the Absorb operation, as it is essential for fully realizing MLA’s potential in practical deployment.

The importance of the absorb operation is further illustrated in Palu [1]. In their Figure 5, they show that removing RoPE allows the absorb operation to achieve a 6.17× speedup at 64K context. However, when RoPE is retained, absorb must be disabled, and the speedup drops to 2.91×. Unfortunately, Palu does not support RoPE decoupling and therefore cannot apply the absorb operation to popular models such as LLaMA, Qwen, and Mistral.

Table 3. Comparison of the impact of num_kv_heads, head_dim, and parameters on inference speed for GQA and MLA. To eliminate the influence of MLP layers on the results, we conducted the speed test using models composed of 32 attention layers only.

Q3.1. Why compare capability under equal KV Cache rather than equal parameter count or head dimension?

A3.1. The inference speed of LLMs is primarily affected by two factors: computational intensity and memory access [7]. As the KV cache length increases during inference, modern hardware primarily hits a memory bottleneck rather than compute bottleneck. For example, in Table 3, GQA models with the same head dimension where one doubles its KV cache size show drastic change in inference speed (also about 2×). In contrast, using the same KV cache size but doubling the head dimensions leads to marginal speed changes. For MLA models, they have the same parameter count before and after absorption, but the difference in KV cache size results in a 16.7× gap in inference speed. Although MLA has more parameters than GQA, its reduced KV cache size leads to faster inference. Therefore, comparing models under equal KV cache size—the factor most directly affecting inference speed—is both practical and meaningful for analyzing expressive capacity.

Q3.2. There are many ways to handle the KV cache—why focus on MLA?

A3.2. As discussed in Related Work, existing KV cache optimization techniques are orthogonal to ours: 1) KV cache pruning reduces the number of tokens stored; 2) KV cache quantization reduces the bit-width per value. In contrast, TransMLA focuses on per-token dimensionality reduction, and thus complements other techniques rather than replacing them.

Q3.3. Is MLA really better than GQA?

A3.3. Theoretically, we prove that MLA always has greater expressive power than GQA under the same KV cache size. Empirically, several SOTA open-source models using 576-dimensional MLA—such as DeepSeek-V3 and Kimi-K2—achieve performance comparable to or even surpassing that of GQA models using 2048-dimensional KV cache, such as LLaMA-4-Maverick-17B-128E, Qwen3-Coder-480B-A35B, Mixtral-8x22B, and Gemma-3-27B. This provides strong evidence that MLA can effectively replace GQA as a core component for high-performance, inference-efficient foundation models.

Q4. Why do GQA models still dominate, and is such a complex conversion to MLA really necessary?

A4. Pretraining a foundation model is a massive system engineering task that requires months of continuous experimentation and iteration. Leading model developers have already invested heavily in GQA-based architectures and accumulated significant optimization experience for both training and inference. Switching to a new architecture would demand substantial system reengineering and delays due to tasks like value alignment before public release, not to mention that DeepSeek R1’s release was only 6 months ago. Despite these challenges, we are already seeing a shift—Kimi-K2, a state-of-the-art agentic model, has adopted MLA. Our theoretical work provides a compelling reason for major model developers to consider this transition. Moreover, the GQA-to-MLA conversion method proposed in this paper allows them to recycle existing pretrained models, avoiding the need to train from scratch and thereby reducing the carbon footprint.

[1] Palu: Compressing KV-Cache with Low-Rank Projection (ICLR 2025)

[2] Towards Economical Inference: Enabling DeepSeek's Multi-Head Latent Attention in Any Transformer-based LLMs (ACL 2025)

[3] X-EcoMLA: Upcycling Pre-Trained Attention into MLA for Efficient and Extreme KV Compression (COLM 2025)

[4] Hardware-Efficient Attention for Fast Decoding

[5] Multi-head Temporal Latent Attention

[6] Grouped-head latenT Attention

[7] Roofline: an insightful visual performance model for multicore architectures.

Thank you for supporting the innovation of converting GQA to MLA, and for recognizing our paper’s Quality, Clarity, Significance, and Originality. In response to your new question, we reply as follows:

Q1. TransMLA is too costly and complex, and many alternative KV-cache optimization methods exist.

A1. The rigorous theoretical guarantees in TransMLA (including the expressivity proof and the proof of rotation invariance under RoPE) may give the impression that the method is costly and complex. Setting those proofs aside, TransMLA ultimately just applies PCA twice—first to decouple RoRoPE, and second to compress the KV cache. The PCA is widely used across various KV-cache compression methods (Palu, ICLR 2025; RazorAttention, ICLR 2025; MatryoshkaKV, ICLR 2025; MHA2MLA, ACL 2025; CLOVER, ICML 2025; Loki, NeurIPS 2024). Therefore, TransMLA itself is not costly or complex. Moreover, as we have already noted, other KV-cache compression techniques such as sparse attention and KV-cache quantization are compatible with TransMLA.

Q2. MLA only helps when training memory is limited. If memory is unconstrained—at equal FLOPs or in an overtrained regime—MLA may not clearly outperform GQA due to low-rank issues.

A2. Pretraining LLMs is extremely expensive. For example: DeepSeek V3 pretraining on 2,788K H800 GPU-hours cost 5.576M dollars, and Llama 3.1 405B pretraining on 30.84M H100 GPU-hours cost 61.68M dollars. Therefore, assuming unconstrained training memory is meaningless. Tri Dao, the author of FlashAttention and Mamba, points out: “LLM decoding is bottlenecked for large batches and long contexts by loading the key-value (KV) cache from high- bandwidth memory, which inflates per-token latency, while the sequential nature of decoding limits parallelism.” [1]

The number of parameters does not directly determine a model’s speed; it influences speed indirectly through floating-point operations (FLOPs) and memory operations. However, “Modern GPUs devote considerably more silicon to computation than memory bandwidth, a trend that intensified in 2017. Hardware FLOPs have scaled by∼3×every two years, while the HBM memory bandwidth increases by∼1.6×over the same period.” [1]

According to the roofline model [2], Arithmetic Intensity = FLOPs / memory operations. Take the NVIDIA H800 SXM5 GPU as an example: it has a peak memory bandwidth of 3.35 TB/s and a practical peak of 865 TFLOPs. When Arithmetic Intensity > 256, the kernel is compute-bound; otherwise, it is memory-bound. For GQA, $\text{Arithmetic Intensity}=\frac{2 h_q s_{kv} d}{2 h_q d + 2 h_{kv} s_{kv} d}\approx \frac{h_q}{k_{kv}}$. For LLaMA-3-70B, $\frac{h_q}{k_{kv}} = \frac{64}{8} = 8$. Thus it is memory-bound, meaning the data fed into the GPU cannot keep up with the GPU’s processing speed, resulting in low GPU utilization. For MLA, $\text{Arithmetic Intensity} = h_q \cdot \frac{d_k + d_v}{d_k} \approx 2 h_q$. For DeepSeek V3, $2 h_q = 256$, which balances compute and memory throughput and therefore lets the GPU achieve high utilization.

This is the theoretical basis for why we emphasize the importance of the KV cache.

Q3. The authors’ design and contributions are only meaningful in the conversion to MLA. After MLA’s introduction, most models still train from scratch using GQA.

A3. Which of GQA or MLA will ultimately dominate large-scale model design remains an open question. This paper argues—on theoretical grounds—that MLA has an expressivity advantage over GQA. Intuitively, each attention head needs to retain some repeated information; under GQA, every group must keep its own copy of that information, whereas under MLA all heads share a single KV cache. Consequently, to preserve the same information, MLA requires a lower rank. Our aim is to draw broader attention to MLA and, additionally, to provide a low-cost pathway for migrating existing models to MLA.

Thank you again for your thoughtful feedback; we would be pleased to continue the discussion as needed. If our updated responses address your concerns, we would be grateful if you could consider raising your score.

[1] Hardware-Efficient Attention for Fast Decoding

[2] Roofline: an insightful visual performance model for multicore architectures.

Dear Reviewer,

We recognize that the transition from GQA to MLA may feel strange to you—here’s one way to understand about it. There are many techniques for compressing the dimensionality of the KV cache that are orthogonal to (i) reducing the number of tokens and (ii) lowering numerical precision via quantization. Building on these dimensionality-compression methods, we further enable MLA’s matrix-absorption capability. This design allows TransMLA to reuse FlashMLA’s inference-optimized kernels, enabling not only easier and faster deployment across DeepSeek-compatible environments compared to approaches that require custom kernels, but also delivering better acceleration performance.

Dear Reviewer i2PF,

Thank you very much for your thoughtful feedback and for raising your score. We sincerely appreciate your input, which has been invaluable in improving our work.

We will ensure that all of your suggestions are carefully considered and incorporated into the camera-ready version of the paper.

Once again, thank you for your support and encouragement.

Best regards