Question: Applicability of TransMLA to DLM

[amandpkr]

Hi Authors,

Thank you for the amazing work and for open-sourcing this project!

I had a quick question — do you think TransMLA could be applied to Diffusion Language Models (DL models)?
I’m currently exploring ways to reduce the K–V cache memory during inference in diffusion-based language or vision–language models.

Since TransMLA provides a theoretical and practical framework for converting GQA-based architectures into MLA with compressed KV caches, I was wondering whether a similar idea could be used for the iterative denoising steps in diffusion models.

Would love to hear your thoughts on whether TransMLA’s low-rank latent compression or RoRoPE decoupling could extend to diffusion-style attention or cross-attention blocks.

Thanks again for this great contribution!

[fxmeng]

Your idea is great — using TransMLA to compress the KV-Cache as initialization for DLM is a very suitable approach. However, there are a few points to keep in mind:
1. For block diffusion models (BDM), TransMLA is a great method for saving KV-Cache.
2. For standard discrete diffusion models, the KV-Cache requires some approximation to be reused. Therefore, TransMLA may have the potential to outperform directly using GQA’s cache (since it retains only the principal components and removes a lot of noise), but it might also lead to reduced approximation quality due to the smaller KV-Cache.

Hence, I suggest you directly port TransMLA to BDM, or compare the performance of TransMLA and GQA cache on DLM.
Looking forward to hearing about your progress — good luck!