Support Qwen3, e.g Qwen3-32B(GQA) can be tranasformed to MLA

[xueliangyang-oeuler]

I want to use transMLA to convert qwen3, so that I can use flashMLA to improve the inference speed of vllm or sglang on qwen3.

[xueliangyang-oeuler]

@yaozengwei @MFX12138 @fxmeng I need some help.

[GiorgioZacharo]

I am also trying to use transMLA for qwen3-8B, please let us know if there is any update and plan to support it. Thanks!

[fxmeng]

Thank you for your interest in the TransMLA. Since Qwen-3 inserts a head-wise RMSNorm between q_proj and k_proj, TransMLA is currently unable, in principle, to perform RoRoPE and KV cache compression across the head-wise RMSNorm (it would work if a standard RMSNorm were used). This remains an issue to be addressed in the future, and we welcome the community to explore potential solutions together.

[xueliangyang-oeuler]

@fxmeng Thank you for your reply. Any there any plans for Qwen-3. Can we make a RFC for this future work?

[fxmeng]

@fxmeng Thank you for your reply. Any there any plans for Qwen-3. Can we make a RFC for this future work?

I don’t have a concrete solution yet, but I agree that preparing an RFC would be helpful. I’m happy to contribute to the discussion and see what ideas others might bring.

[amandpkr]

@fxmeng I have a generally doubt about the normalization of RMSNorm, does applying RMS normalization after RoPE effect the performance ?

[xueliangyang-oeuler]

I am also trying to use transMLA for qwen3-8B, please let us know if there is any update and plan to support it. Thanks!我也在尝试使用 transMLA 来支持 qwen3-8B，请问是否有任何更新或计划支持该版本？谢谢！

Hi, are there any breakthrough for qwen3-8b? Can you share some ideas or work on this? @GiorgioZacharo

[xueliangyang-oeuler]

@fxmeng Thank you for your reply. Any there any plans for Qwen-3. Can we make a RFC for this future work?@fxmeng 谢谢你的回复。Qwen-3 有什么计划吗？我们可以为此未来的工作写一份 RFC 吗？

I don’t have a concrete solution yet, but I agree that preparing an RFC would be helpful. I’m happy to contribute to the discussion and see what ideas others might bring.我目前还没有具体的解决方案，但我同意编写一份 RFC（征求意见稿）会很有帮助。我很乐意参与讨论，看看其他人会提出什么想法。

Can you issue a RFC for qwen3 models supporting transformed to MLA? We can make steps for this work. We can join in.

[nagarajankarthik]

@fxmeng ,

We are trying to understand the limitations of the TransMLA framework with regard to the Qwen3 models. Let us consider this simple example:

• Size of hidden dimension per key-value head = $d$ = 4
• Number of key and value heads = $g$ = 2.
• Number of query heads = $h$.
• freqfold parameter = $\phi$ = 2

In the absence of the RMSNorm operation, the dot product of the query vector from the first query head and the key vector for a single token can be written as

$$ [q_0, q_1, q_2, q_3, q_4, q_5, q_6, q_7]^{R^T} \cdot [k_0, k_1, k_2, k_3, k_4, k_5, k_6, k_7]^R $$

Upon grouping real and imaginary components, one obtains

$$ ([q_0, q_2, q_4, q_6]; [q_1, q_3, q_5, q_7])^{R^T} \cdot ([k_0, k_2, k_4, k_6]; [k_1, k_3, k_5, k_7])^R $$

As per Equation 19 and Appendix B of your paper, the above is equivalent to

$$ (U[q_0, q_2, q_4, q_6]; U[q_1, q_3, q_5, q_7])^{R^T} \cdot (U[k_0, k_2, k_4, k_6]; U[k_1, k_3, k_5, k_7])^R $$

, where $U$ is the rotation matrix obtained from the model query outputs on a calibration dataset.

If one performs a single RMSNorm for queries across all query heads without introducing an additional per-channel scaling factor, and does likewise for the keys, the expression becomes

$$ \frac{1}{q_n}(U[q_0, q_2, q_4, q_6]; U[q_1, q_3, q_5, q_7])^{R^T} \cdot \frac{1}{k_n}(U[k_0, k_2, k_4, k_6]; U[k_1, k_3, k_5, k_7])^R $$

where $q_n = \sqrt{\sum_{t = 0}^{hd - 1} q_t^2}$ and $k_n = \sqrt{\sum_{t = 0}^{gd - 1} k_t^2}$.

The ability to extract the factors $\frac{1}{q_n}$ and $\frac{1}{k_n}$ outside the matrix multiplication with $U$ indicates that the RMSNorm can be performed after the rotation.

However, if the keys and queries are normalized headwise with a separate scaling factor for each channel, one obtains

$$ (U[\frac{\alpha_0}{q_{n0}} q_0, \frac{\alpha_2}{q_{n0}} q_2, \frac{\alpha_4}{q_{n0}} q_4, \frac{\alpha_6}{q_{n0}} q_6]; U[\frac{\alpha_1}{q_{n0}} q_1, \frac{\alpha_3}{q_{n0}} q_3, \frac{\alpha_5}{q_{n0}} q_5, \frac{\alpha_7}{q_{n0}} q_7])^R \cdot ( U[ \frac{\beta_0}{k_{n0}} k_0, \frac{\beta_2}{k_{n0}} k_2, \frac{\beta_0}{k_{n1}} k_4, \frac{\beta_2}{k_{n1}} k_6]; U[\frac{\beta_1}{k_{n0}} k_1, \frac{\beta_3}{k_{n0}} k_3, \frac{\beta_1}{k_{n1}} k_5, \frac{\beta_3}{k_{n1}} k_7])^R $$

where $q_{n0} = \sqrt{\sum_{t = 0}^{gd - 1} q_t^2}$, $k_{ni} = \sqrt{\sum_{t = id}^{id + d - 1} k_t^2}$ and $\alpha_j$ and $\beta_j$ are the scaling factors for channel $j$.

In this case, the RMSNorm of queries and keys must be performed before the rotation operation used in RoRoPE as the RMSNorm coefficients are different for each element of the query and key vectors. Based on my understanding, this implies that $U$ cannot be multiplied together with $W^{UK}$ and $W^K$ as mentioned on page 7 of the TransMLA paper. However, it is not possible to optimize $U$ independently of $W^{UK}$ and $W^K$ due to the requirement that $U^TU = I$.

Could you please clarify if this is the reason that prevents TransMLA from being applied to Qwen3 in principle?

Thanks.