AMLA: MUL by ADD in FlashAttention Rescaling
On the map: Efficient LLM inference and serving → KV cache optimization
Abstract (arXiv)
Multi-head Latent Attention (MLA) significantly reduces KVCache memory usage in Large Language Models while introducing substantial computational overhead and intermediate variable expansion. This poses challenges for efficient hardware implementation -- especially during the decode phase. This paper introduces Ascend MLA (AMLA), a high-performance kernel specifically optimized for Huawei's Ascend NPUs. AMLA is built on two core innovations: (1) A novel FlashAttention-based algorithm that replaces floating-point multiplications with integer additions for output block rescaling, leveraging binary correspondence between FP32 and INT32 representations; (2) A Preload Pipeline strategy with hierarchical tiling that maximizes FLOPS utilization: the Preload Pipeline achieves Cube-bound performance, while hierarchical tiling overlaps data movement and computation within the Cube core. Experiments show that on Ascend 910 NPUs (integrated in CloudMatrix384), AMLA achieves up to 614 TFLOPS, reaching 86.8% of the theoretical maximum FLOPS, outperforming the state-of-the-art open-source FlashMLA implementation, whose FLOPS utilization is up to 66.7% on NVIDIA H800 SXM5. The AMLA kernel has been integrated into Huawei's CANN and will be released soon.
Key facts. Cited by 1 paper in the map (data through 14 Oct 2026). Placed in KV cache optimization, one of 12 directions in Efficient LLM inference and serving. Builds on 17 of its 21 references that are in the map. First posted on arXiv 24 Sep 2025.
Builds on
17 of 21 references.
Closest ancestors · 1
Same direction, newest first
Other references · 16
- 2017Attention Is All You Needcited by 60,361
- 2020Language Models are Few-Shot Learnerscited by 32,419
- 2023GPT-4 Technical Reportcited by 18,396
- 2022ReAct: Synergizing Reasoning and Acting in Language Modelscited by 8,695
- 2025DeepSeek-R1: Incentivizing Reasoning Capability in LLMs via Reinforcement Learningcited by 6,330
- 2022FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awarenesscited by 3,673
- 2023Qwen Technical Reportcited by 3,344
- 2023FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioningcited by 2,598
- 2023GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpointscited by 1,346
- 2024DeepSeek-V2: A Strong, Economical, and Efficient Mixture-of-Experts Language Modelcited by 1,135
- 2019Fast Transformer Decoding: One Write-Head is All You Needcited by 717
- 2018Online normalizer calculation for softmaxcited by 162
4 more
- 2024Large Language Models Meet NLP: A Surveycited by 103
- 2024A Survey of Large Language Models for Financial Applications: Progress, Prospects and Challengescited by 90
- 2024Large Language Models for Manufacturingcited by 24
- 2025Serving Large Language Models on Huawei CloudMatrix384cited by 24
Built on by
Citing papers per year: 2026: 1.
1 paper.
Similar, not cited
Close in wording, no citation either way
- 2026MLAfication: Efficiently Enabling High-Fidelity MLA for Modern LLMs
- 2026W4A16 Mixed-Precision Matrix Multiplication on Decoupled Architecture: Kernel Design and Memory Bottleneck Analysis for Ascend NPUs
- 2025Hardware-Centric Analysis of DeepSeek's Multi-Head Latent Attention
- 2025Multi-head Temporal Latent Attention
- 2026Vectorization Of Narrow Matrix Multiplication for Ascend AI Inference Acceleration
- 2024INT-FlashAttention: Enabling Flash Attention for INT8 Quantization
- 2024FlashAttention-3: Fast and Accurate Attention with Asynchrony and Low-precision
- 2026FlashAttention for Scalable Vector Architectures