FlashMLA-ETAP: Efficient Transpose Attention Pipeline for Accelerating MLA Inference on NVIDIA H20 GPUs

2025-05-13 · Pengcuo Dege, Qiuming Luo, Rui Mao, Chang Kong

Efficient inference of Multi-Head Latent Attention (MLA) is challenged by deploying the DeepSeek-R1 671B model on a single Multi-GPU server. This paper introduces FlashMLA-ETAP, a novel framework that enhances MLA inference for the single-instance deployment scenario on NVIDIA H20 GPUs. We propose the Efficient Transpose Attention Pipeline (ETAP), which reconfigures attention computation through transposition to align the KV context length with the $M$-dimension in WGMMA operations, significantly reducing redundant computations. 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, while maintaining numerical stability with a 15.2x lower RMSE ($1.25 \times 10^{-5}$) than FlashAttention-3. Furthermore, ETAP's design enables seamless integration into frameworks like FlashAttention-3 and FlashInfer, supported by a detailed theoretical analysis. Our work addresses a critical gap in resource-constrained inference, offering a scalable solution for mid-tier GPUs and paving the way for broader adoption in hardware-aware optimization. Code is available at https://github.com/pengcuo/FlashMLA-ETAP.

📄 PDF Abstract BibTeX arXiv:2506.01969

Code (1)

pengcuo/flashmla-etap 공식 구현 pytorch

Tasks

GPU

Methods 이 논문이 사용한 방법론

Softmax The Softmax output function transforms a previous layer's output into a vector of probabilities. It is commonly used for multiclass classification. Given an input vector $x$…

Attention 설명 없음

ALIGN In the ALIGN method, visual and language representations are jointly trained from noisy image alt-text data. The image and text encoders are learned via contrastive loss…

AMLA: MUL by ADD in FlashAttention Rescaling

2025-09-24 · Qichen Liao, Chengqiu Hu, Fangzheng Miao, Bao Li 외 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 ef…

Finding Challenging Metaphors that Confuse Pretrained Language Models

2024-01-29 · Yucheng Li, Frank Guerin, Chenghua Lin

Metaphors are considered to pose challenges for a wide spectrum of NLP tasks. This gives rise to the area of computational metaphor processing. However, it remains unclear what types of metaphors challenge current state-…

Machine Translation

Metaphor Detection using Deep Contextualized Word Embeddings

2020-09-26 · Shashwat Aggarwal, Ramesh Singh

Metaphors are ubiquitous in natural language, and their detection plays an essential role in many natural language processing tasks, such as language understanding, sentiment analysis, etc. Most existing approaches for m…

Sentiment Analysis Word Embeddings

TransPose: Keypoint Localization via Transformer

2020-12-28 · ICCV 2021 10 · Sen yang, Zhibin Quan, Mu Nie, Wankou Yang

While CNN-based models have made remarkable progress on human pose estimation, what spatial dependencies they capture to localize keypoints remains unclear. In this work, we propose a model called \textbf{TransPose}, whi…

Keypoint Detection Multi-Person Pose Estimation Pose Estimation

The Mind's Eye: A Multi-Faceted Reward Framework for Guiding Visual Metaphor Generation

2025-08-26 · Girish A. Koushik, Fatemeh Nazarieh, Katherine Birch, Shenbin Qian 외 arxiv

Visual metaphor generation is a challenging task that aims to generate an image given an input text metaphor. Inherently, it needs language understanding to bind a source concept with a target concept, in a way that pres…