FlashMLA-ETAP: 高效转置Attention流水线加速H20 GPU上的MLA推理

AbstractFlashMLA-ETAP提出了一种针对NVIDIA H20 GPU单实例部署场景优化的Multi-Head Latent Attention (MLA)推理框架。通过引入高效转置Attention流水线(ETAP)，重构attention计算以减少冗余操作，并将KV context长度与WGMMA操作的M维度对齐，充分利用H20硬件特性。在64K序列长度、batch size 16的场景下，相比FlashMLA实现2.78倍加速，相比FlashAttention-3和FlashInfer分别实现5.24倍和4.94倍提升。同时保持数值稳定性，RMSE比FlashMLA低15.2倍。 Key Contributions Efficient Transpose Attention Pipeline (ETAP): 通过转置重构attention计算，减少冗余操作并优化硬...

阅读全文