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

Technical Summary: FlashMLA-ETAP

Problem Statement
The deployment of large-scale models, such as DeepSeek-R1 (671B parameters), on resource-constrained mid-tier GPUs like the NVIDIA H20 presents significant efficiency challenges for Multi-Head Latent Attention (MLA) inference. While high-end GPUs (e.g., H100, H800) possess sufficient compute capability to handle existing optimizations, the NVIDIA H20’s modest 148 TFLOPS FP16 compute capability creates a bottleneck. Specifically, in single-instance deployment scenarios where the model’s 128 attention heads are split across 8 GPUs (resulting in 16 heads per GPU), the per-GPU head count falls below the minimum M-dimension requirement (64) of WarpGroup Matrix-Multiply-Accumulate (WGMMA) instructions. This mismatch necessitates redundant padding, causing compute utilization to drop below 25% during decoding phases characterized by short query lengths and long Key-Value (KV) context lengths. Existing frameworks like FlashAttention-3 and FlashMLA, optimized for high-end hardware, fail to fully utilize the H20’s resources due to these architectural constraints and suboptimal memory access patterns.

Methodology
To address these inefficiencies, the authors propose FlashMLA-ETAP, a framework that integrates the Efficient Transpose Attention Pipeline (ETAP) into the FlashMLA architecture.

• Efficient Transpose Attention Pipeline (ETAP): ETAP reconfigures the standard attention computation by leveraging transposition to align computational dimensions with the H20’s WGMMA constraints. In traditional MLA inference, the computation follows $S = Q \cdot K^T$, $P = \text{softmax}(S)$, and $O = P \cdot V$. ETAP reorients this process to:

  1. $S^T = K \cdot Q^T \in \mathbb{R}^{N \times N_q}$
  2. $P^T = \text{softmax}(S^T) \in \mathbb{R}^{N \times N_q}$
  3. $O = V^T \cdot P^T \in \mathbb{R}^{d \times N_q}$
  4. $O = O^T$

  By treating the long KV context length ($N$) as the M-dimension and the short query length ($N_q$) as the N-dimension in WGMMA operations, ETAP eliminates the need for redundant padding on the query dimension. This alignment exploits the asymmetry between query and KV lengths in autoregressive decoding, significantly reducing computational overhead.

• Integration and Implementation: FlashMLA-ETAP implements ETAP within FlashMLA’s Cooperative Thread Array (CTA) framework. It utilizes producer-consumer synchronization and block-wise processing with a circular shared memory buffer. The algorithm splits computation between two warpgroups: a consumer warpgroup handling attention scores and softmax statistics, and a producer warpgroup managing data loading from High-Bandwidth Memory (HBM) to Shared Memory (SMEM). This design preserves FlashMLA’s low-rank joint compression for KV cache minimization while optimizing the matrix multiplication and softmax stages for the H20.

• Theoretical Generalizability: The paper provides a theoretical analysis demonstrating that ETAP’s transposition strategy is compatible with other attention frameworks, including FlashAttention-3 and FlashInfer. Since these frameworks share the fundamental attention computation pattern, ETAP can be integrated with minimal modifications to their tiling and kernel fusion strategies, potentially extending its benefits to other mid-tier GPU architectures.

Key Contributions

1. Introduction of ETAP: A novel computation mode that uses transposition to reduce redundant calculations caused by WGMMA padding requirements on mid-tier GPUs.
2. Development of FlashMLA-ETAP: An enhanced version of FlashMLA that incorporates ETAP to enable efficient MLA inference on the NVIDIA H20, with a design intended for seamless integration into other frameworks like FlashAttention-3 and FlashInfer.
3. Performance Validation: Experimental results demonstrating significant speedups and improved numerical stability compared to baseline frameworks on the H20 GPU.

Results
Experiments were conducted using the DeepSeek-R1 model on an NVIDIA H20 GPU, comparing FlashMLA-ETAP against FlashMLA, FlashAttention-3, and FlashInfer across sequence lengths from 512 to 64K and batch sizes of 16 and 32.

• Speedup: At a sequence length of 64K and batch size 16, FlashMLA-ETAP achieved a 2.78x speedup over FlashMLA (89 TFLOPS/s vs. 32 TFLOPS/s). It also outperformed FlashAttention-3 by 5.24x and FlashInfer by 4.94x. For batch size 32, it achieved a 2.72x speedup over FlashMLA at 64K.
• Scalability: The performance advantage of ETAP increased with sequence length, highlighting its effectiveness in mitigating padding overhead in long-context scenarios.
• Numerical Stability: FlashMLA-ETAP demonstrated superior numerical accuracy compared to FlashAttention-3. The Root Mean Square Error (RMSE) against a double-precision reference was $1.25 \times 10^{-5}$ for FlashMLA-ETAP, compared to $1.9 \times 10^{-4}$ for FlashAttention-3, representing a 15.2x reduction in error.

Significance
The paper claims that FlashMLA-ETAP addresses a critical gap in resource-constrained inference by offering a scalable solution for mid-tier GPUs. By optimizing for the specific architectural constraints of the NVIDIA H20, the framework enables efficient deployment of large models like DeepSeek-R1 in single-instance scenarios. The authors posit that ETAP’s generalizable design paves the way for broader adoption of hardware-aware optimizations, potentially democratizing efficient inference on mid-tier hardware and supporting future research in computational efficiency for diverse GPU architectures.