Computing attention is the backbone of transformer-based models like large language models. However, the increasing diversity of attention algorithms presents significant challenges for unleashing hardware performance. State-of-the-art variants like FlashAttention target a specific attention algorithm or hardware platform, which fail to generalize to other algorithms and platforms.

We present MetaAttention, a framework that automatically derives the optimal implementation of an attention algorithm given a hardware platform.
Our key insight is that variants of attention can be abstracted into two operations:
relevance scoring and aggregation, complemented by customizable functions and configurations like the input shape.
Based on it, we systematically design a cross-backend attention runtime around these operations that generalizes to variants of attention with customizable operators.
To unleash the hardware performance, we further propose an IntermediateTensor-based search method to find the optimal tiling strategy and the parallelism scheme according to the attention customization and hardware features.
MetaAttention delivers up to a 10.4$\times$ speedup for configurations previously unsupported by state-of-the-art systems.
Additionally, MetaAttention achieves performance comparable to manually-optimized libraries such as FlashMLA while significantly reducing the amount of code required.