Q i M eng-Attention: SOTA Attention Operator is generated by SOTA Attention Algorithm (2025)

Abstract: The attention operator remains a critical performance bottleneck in large language models (LLMs), particularly for long-context scenarios. While FlashAttention is the most widely used and effective GPU-aware acceleration algorithm, it must require time-consuming and hardware-specific manual implementation, limiting adaptability across GPU architectures. Existing LLMs have shown a lot of promise in code generation tasks, but struggle to generate high-performance attention code. The key challenge is it cannot comprehend the complex data flow and computation process of the attention operator and utilize low-level primitive to exploit GPU performance.To address the above challenge, we propose an LLM-friendly Thinking Language (LLM-TL) to help LLMs decouple the generation of high-level optimization logic and low-level implementation on GPU, and enhance LLMs’ understanding of attention operator.Along with a 2-stage reasoning workflow, TL-Code generation and translation, the LLMs can automatically generate FlashAttention implementation on diverse GPUs, establishing a self-optimizing paradigm for generating high-performance attention operators in attention-centric algorithms.Verified on A100, RTX8000, and T4 GPUs, the performance of our methods significantly outshines that of vanilla LLMs, achieving a speed-up of up to 35.16\times .Besides, our method not only surpasses human-optimized libraries (cuDNN and official library) in most scenarios but also extends support to unsupported hardware and data types, reducing development time from months to minutes compared with human experts.

## introduction the attention mechanism is the cornerstone of modern large language models (llms). its time and memory complexity grows quadratically in * corresponding author. sequence length, and attention has emerged as the bottleneck for the runtime and memory resource requirements for llms, particularly in long-context scenarios @xcite . consequently, numerous acceleration algorithms have been proposed to accelerate attention operators, with flashattention being the most effective and widely used @xcite @xcite . with the rise of cost-efficient, high-performance llms like deepseek @xcite , deploying llms across diverse gpu specifications (including legacy generations) has become increasingly prevalent. however, due to the complexity of attention operators and disparities in the architectures and instruction sets across different gpus, the efficient implementation and migration of flashattention algorithms becomes a pressing challenge. exist-ing methods for implementing high-performance attention operators can be divided into two categories: i.e., manual implementation and deep learning compilers. human experts optimize attention operators through fused computation based on their expertise in gpu architecture and algorithm knowledge, thereby effectively enhancing computation efficiency. consequently, these handcrafted implementations are time-consuming and non-portable, failing to support some common gpus and datatypes (e.g., rtx8000 gpu and fp8 datatype). in addition, existing mainstream deep learning compilers (e.g., tvm @xcite , ansor @xcite ) fail to automatically implement flashattention due to their inability to perform optimizations involving fused sequences of multiple complex operations (e.g. gemm). overall, current methods struggle to automatically analyze the attention operator characteristics and gpu architecture to efficiently generate high-performance flashattention implementation on different gpu platforms, which limits the adaptability of llms. motivated by strong code-generation abilities of llms @xcite @xcite @xcite @xcite , it is a promising direction to tackle these issues by utilizing llms automatically implement flashattention on gpus. although llms can create attention implementations by high-level apis like pytorch @xcite , they have trouble generating highperformance flashattention with low-level primitives (e.g., cuda or cute (nvidia/cutlass, 2023)) with crucial optimization techniques like data blocking and fused computation. this significantly hinders the reduction of memory access overheads and the enhancement of computational efficiency, as shown in figure 1(a) .the primary challenges preventing llms from generating highperformance flashattention implementations in one pass are twofold: (1) llms cannot comprehend the complex data flow and multi-step computation of attention, thus failing to generate effective optimization logic; (2) llms are unable to produce optimized implementations based on low-level primitives to fully exploit gpu characteristics. inspired by human experts' decoupling of optimization and implementation, we propose an llmfriendly thinking language (llm-tl) to address the above challenge by helping llms decouple the generation of abstract optimization logic and low-level code implementation on gpu. specifically, llm-tl encompasses two types of statements and their requisite parameters: memory access and computation, enabling llms to comprehend the data flow and computation processes of the attention operator on gpu architecture from the vantage point of high-level abstract semantics, as shown in figure 1 the contributions of this paper are as follows: ## evaluation 4.1 experiment setup to validate the performance of the attention operator implementation generated by llms through llm-tl, we conduct comprehensive evaluation across different gpu platforms, distinct llms, and attention variants. additional experiments are shown in appendix a due to space limitations. gpu platforms. llm-tl is test on gpus with different architectures, including nvidia ampere (a100) and turing (rtx8000, t4) architecture. llms. the proposed framework is validated on four sota llms, including the closed-source models of @xcite comparison baselines. implementation generated by our approach is compared with flashattention official library (v2.7.3 on a100 and v1.0.9 on rtx8000, as flashattention v2 is not available on turing architecture), nvidia official deep learning library cudnn @xcite , flex-attention official library @xcite , benchmark setting. we followed the settings in the flashattention @xcite , the sequence length varies from 512, 1k, ..., 16k, and the batch size is adjusted such that the total number of tokens remains 16k. we set the hidden dimension to 2048, and the head dimension to be either 64 or 128 (i.e., 32 heads or 16 heads). the flops are calculated as: 4 * seqlen 2 * head dimension * number of heads. for mla, we utilize the dimensions specified in deepseek-v3, which consist of an embedding dimension of 128 and a rope dimension of 64. we also test attention variants with and without causal masks.