AMLA: MUL by ADD in FlashAttention Rescaling

• URL: http://arxiv.org/abs/2509.25224v1
• Date: Wed, 24 Sep 2025 10:04:59 GMT
• Title: AMLA: MUL by ADD in FlashAttention Rescaling
• Authors: Qichen Liao, Chengqiu Hu, Fangzheng Miao, Bao Li, Yiyang Liu, Junlong Lyu, Lirui Jiang, Jun Wang, Lingchao Zheng, Jun Li, Yuwei Fan,
• Abstract summary: Multi-head Latent Attention (MLA) significantly reduces KVCache memory usage in Large Language Models.<n>This paper introduces Ascend MLA, a high-performance kernel specifically optimized for Huawei's Ascend NPUs.
• Score: 9.20326267673456
• License: http://arxiv.org/licenses/nonexclusive-distrib/1.0/
• Abstract: 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.

Related papers

• LFM2 Technical Report [87.58431408281973]
  We present LFM2, a family of Liquid Foundation Models designed for efficient on-device deployment and strong task capabilities.<n>The LFM2 family covers 350M-8.3B parameters, including dense models (350M, 700M, 1.2B, 2.6B) and a mixture-of-experts variant (8.3B total, 1.5B active)<n>We build multimodal and retrieval variants: LFM2-VL for vision-latency tasks, LFM2-Audio for speech, and LFM2-ColBERT for retrieval.
  arXiv  Detail & Related papers  (2025-11-28T17:56:35Z)
• RISC-V Based TinyML Accelerator for Depthwise Separable Convolutions in Edge AI [1.1816942730023885]
  This paper introduces a novel hardware accelerator architecture that utilizes a fused pixel-wise dataflow.<n>It computes a single output pixel to completion across all stages-expansion, depthwise convolution, and projection-by streaming data.<n>It achieves a speedup of up to 59.3x over the baseline software execution on the RISC-V core.
  arXiv  Detail & Related papers  (2025-11-26T10:01:31Z)
• Eliminating Multi-GPU Performance Taxes: A Systems Approach to Efficient Distributed LLMs [61.953548065938385]
  We introduce the ''Three Taxes'' (Bulk Synchronous, Inter- Kernel Data Locality, and Kernel Launch Overhead) as an analytical framework.<n>We propose moving beyond the rigid BSP model to address key inefficiencies in distributed GPU execution.<n>We observe a 10-20% speedup in end-to-end latency over BSP-based approaches.
  arXiv  Detail & Related papers  (2025-11-04T01:15:44Z)
• xLLM Technical Report [57.13120905321185]
  We introduce xLLM, an intelligent and efficient Large Language Model (LLM) inference framework.<n>xLLM builds a novel decoupled service-engine architecture.<n>xLLM-Engine co-optimizes system and algorithm designs to fully saturate computing resources.
  arXiv  Detail & Related papers  (2025-10-16T13:53:47Z)
• dInfer: An Efficient Inference Framework for Diffusion Language Models [54.80918957287927]
  Diffusion-based large language models (dLLMs) have emerged as a promising alternative to autoregressive (AR) LLMs.<n>We present dInfer, an efficient and framework for dLLM inference.
  arXiv  Detail & Related papers  (2025-10-09T16:19:42Z)
• Fast-dLLM v2: Efficient Block-Diffusion LLM [64.38006546510337]
  Fast-dLLM v2 is a block diffusion language model that adapts pretrained AR models into dLLMs for parallel text generation.<n>This represents a 500x reduction in training data compared to full-attention diffusion LLMs such as Dream (580B tokens)
  arXiv  Detail & Related papers  (2025-09-30T14:40:18Z)
• Serving Large Language Models on Huawei CloudMatrix384 [28.88558053380112]
  Traditional AI clusters face limitations in compute intensity, memory bandwidth, inter-chip communication, and latency.<n>This paper introduces Huawei CloudMatrix, a next-generation AI architecture realized in the production-grade CloudMatrix384 supernode.<n>It integrates 384 Ascend 910 NPUs and 192 Kunpeng CPUs interconnected via an ultra-high-bandwidth Unified Bus (UB) network, enabling direct all-to-all communication and dynamic pooling of resources.
  arXiv  Detail & Related papers  (2025-06-15T03:41:34Z)
• FlashDMoE: Fast Distributed MoE in a Single Kernel [2.246222223318928]
  FlashDMoE is a fully GPU-resident MoE operator that fuses expert computation and inter-GPU communication into a emphsingle persistent GPU kernel.<n>We show that FlashDMoE achieves up to textbf9$times$ higher GPU utilization, textbf6$times$ lower latency, textbf5.7$times$ higher throughput, and textbf4$times$ better overlap efficiency compared to state-of-the-art baselines.
  arXiv  Detail & Related papers  (2025-06-05T06:29:14Z)
• FlashMLA-ETAP: Efficient Transpose Attention Pipeline for Accelerating MLA Inference on NVIDIA H20 GPUs [2.9406946721643092]
  This paper introduces FlashMLA-ETAP, a novel framework that enhances MLA inference for the single-instance deployment scenario.<n>ETAP reconfigures attention computation through transposition to align the KV context length with the (M)-dimension in WGMMA operations.<n>FlashMLA-ETAP achieves a 2.78x speedup over FlashMLA at 64K sequence length (batch size 16), with 5.24x and 4.94x improvements over FlashAttention-3 and FlashInfer, respectively.
  arXiv  Detail & Related papers  (2025-05-13T17:45:34Z)
• MobileMamba: Lightweight Multi-Receptive Visual Mamba Network [51.33486891724516]
  Previous research on lightweight models has primarily focused on CNNs and Transformer-based designs. We propose the MobileMamba framework, which balances efficiency and performance. MobileMamba achieves up to 83.6% on Top-1, surpassing existing state-of-the-art methods.
  arXiv  Detail & Related papers  (2024-11-24T18:01:05Z)
• EPS-MoE: Expert Pipeline Scheduler for Cost-Efficient MoE Inference [49.94169109038806]
  This paper introduces EPS-MoE, a novel expert pipeline scheduler for MoE that surpasses the existing parallelism schemes.<n>Our results demonstrate at most 52.4% improvement in prefill throughput compared to existing parallel inference methods.
  arXiv  Detail & Related papers  (2024-10-16T05:17:49Z)
• vTensor: Flexible Virtual Tensor Management for Efficient LLM Serving [53.972175896814505]
  Large Language Models (LLMs) are widely used across various domains, processing millions of daily requests. Large Language Models (LLMs) are widely used across various domains, processing millions of daily requests.
  arXiv  Detail & Related papers  (2024-07-22T14:37:58Z)
• DISTFLASHATTN: Distributed Memory-efficient Attention for Long-context LLMs Training [82.06732962485754]
  FlashAttention effectively reduces the quadratic peak memory usage to linear in training transformer-based large language models (LLMs) on a single GPU. We introduce DISTFLASHATTN, a memory-efficient attention mechanism optimized for long-context LLMs training. It achieves 1.67x and 1.26 - 1.88x speedup compared to recent Ring Attention and DeepSpeed-Ulysses.
  arXiv  Detail & Related papers  (2023-10-05T03:47:57Z)
• FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioning [11.508362885430133]
  We exploit the asymmetric GPU memory hierarchy to bring significant memory saving and runtime speedup. FlashAttention is still not nearly as fast as optimized matrix-multiply (GEMM) operations, reaching only 25-40% of the theoretical maximum FLOPs/s. We propose FlashAttention-2, with better work partitioning to address these issues.
  arXiv  Detail & Related papers  (2023-07-17T17:50:36Z)
• Edge-MoE: Memory-Efficient Multi-Task Vision Transformer Architecture with Task-level Sparsity via Mixture-of-Experts [60.1586169973792]
  M$3$ViT is the latest multi-task ViT model that introduces mixture-of-experts (MoE) MoE achieves better accuracy and over 80% reduction computation but leaves challenges for efficient deployment on FPGA. Our work, dubbed Edge-MoE, solves the challenges to introduce the first end-to-end FPGA accelerator for multi-task ViT with a collection of architectural innovations.
  arXiv  Detail & Related papers  (2023-05-30T02:24:03Z)

This list is automatically generated from the titles and abstracts of the papers in this site.

This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.