Zebra-Llama and X-EcoMLA: New Paradigms for Efficient Large Model Inference Proposed by AMD

The AMD research team has proposed two technologies, Zebra-Llama and X-EcoMLA. Through a hybrid architecture (combining State Space Model (SSM) and Multi-Head Latent Attention (MLA)) and KV cache compression, they achieve a significant improvement in large model inference efficiency using only billions of training tokens (far fewer than the trillions required for scratch pre-training). The KV cache compression rate reaches up to over 97%, and no retraining of the model is needed, providing an efficient upgrade path for already deployed Large Language Models (LLMs).

With the widespread deployment of LLMs in various scenarios, inference efficiency has become a key bottleneck for their popularization. Traditional Transformer architectures need to store a large amount of Key-Value (KV) cache during inference, and memory usage grows linearly with sequence length, limiting long-context processing capabilities. The industry currently has two solution approaches: one uses SSM (e.g., the Mamba series) to replace the attention mechanism with recursive states; the other uses MLA technology to reduce KV cache via low-rank compression. However, both require scratch pre-training, which is extremely costly.

The core innovation of X-EcoMLA is a "post-training distillation" method that can upgrade pre-trained Transformer models to efficient MLA variants without scratch training. This technology uses the "dark knowledge" of the original model for lightweight adaptation, achieving extreme KV cache compression while maintaining performance. Experimental data: For the Llama3.2-1B-Instruct baseline model, the average score remains unchanged after 6.4x KV cache compression (requiring only 3.6B training tokens and 70 AMD MI300 GPU hours); the 10.6x compression version loses less than 0.1% of the average score (using 7B tokens and 140 GPU hours).

Zebra-Llama adopts a hybrid architecture combining SSM and MLA layers, efficiently transferring knowledge from pre-trained Transformer models through fine initialization and post-training processes. It builds a model family of three scales: 1B, 3B, and 8B, which can be completed with only 7-11B training tokens (far fewer than the trillions required for pre-training). In terms of KV cache compression: the 1B model is reduced to 3.9% of the original, the 3B model to 2%, and the 8B model to 2.73%, while retaining over 100%, 100%, and 97% of zero-shot performance respectively.

Compared with similar methods like MambaInLLaMA and Minitron, Zebra-Llama has significant advantages: In training efficiency, Zebra-Llama-8B (using an 8B teacher model) achieves 7% higher few-shot accuracy than Minitron-8B (using a 15B teacher model), with 8x fewer training tokens and over 12x smaller KV cache; In inference throughput, it is 2.6-3.8x higher than MambaInLlama under a 32K context length; In memory efficiency, the KV cache compression rate reaches up to over 97%.

The core value of the two technologies lies in providing a "model upgrade" rather than a "retraining" path. Deployed LLM applications can improve inference efficiency through post-training adaptation without bearing the huge cost of scratch training. Application scenarios include: edge device deployment (low memory usage), long-context processing (ultra-long documents/video sequences), and cost optimization (accelerating commercial scenario implementation).

AMD has open-sourced the relevant code on GitHub, providing a complete training and inference workflow. Related papers are published on arXiv (X-EcoMLA: arXiv:2503.11132; Zebra-Llama: arXiv:2505.17272). The research team stated that they will release model checkpoints after the papers are accepted. This direction is an important progress for efficient large model inference, promoting the popularization of LLMs in wider scenarios.