AttnFuse: A Composable DSL for Compiling Attentions to Fused GPU Kernels

# AttnFuse: A Composable DSL for Compiling Attentions to Fused GPU Kernels
Source: [https://arxiv.org/html/2609.13612](https://arxiv.org/html/2609.13612)
Varun Kumar DasojuAffiliation:Department of Computer ScienceAffiliation:University of Wisconsin–MilwaukeeEmail:[dasoju@uwm\.edu](mailto:)Tian ZhaoAffiliation:Department of Computer ScienceAffiliation:University of Wisconsin–MilwaukeeEmail:[tzhao@uwm\.edu](mailto:)

###### Abstract

Modern AI systems are built on the Transformer architecture, whose core operation, attention, accounts for the majority of computation and memory cost\. Researchers continually propose new attention variants to improve quality, efficiency, or context length, but each variant currently requires expert\-written GPU code to run at usable speeds\. PyTorch’s recent flex\_attention lets researchers describe custom attention patterns in Python and compile them to fused kernels, but its design is limited to modifications applied after the central matrix multiplication, excluding Rotary Position Embedding \(RoPE\), the positional encoding used by every major LLM\.

We introduce AttnFuse, a small DSL for attention that makes pre\-multiplication transformations like RoPE first\-class operations\. Researchers compose ten high\-level building blocks to describe a variant, and AttnFuse’s compiler emits a single fused GPU kernel for the entire computation\. On an RTX 3090, AttnFuse achieves a 2\.10×\\timesspeedup over flex\_attention on the RoPE\+causal pattern\. On an H100, it runs a full Llama\-3\-8B training step within 5% of PyTorch’s hand\-tuned backend\. Our investigation reveals the Rotation Calculus: whether to fuse RoPE or apply it separately depends on the GPU’s compute\-to\-bandwidth ratio, with a derived crossover that matches measurement\. AttnFuse demonstrates that a small, attention\-specific compiler can close the gap between flexible research code and production kernels\.

## 1Introduction

#### The attention bottleneck\.

At the heart of every modern language model lies a single operation called*attention*\. Given a sequence of tokens, attention computes which prior tokens are most relevant to each new token, then aggregates their information\. This operation accounts for the majority of both computation and memory cost in models like ChatGPT, Claude, and LLaMA\. Making attention faster is therefore one of the central problems in AI systems research\.

The challenge is that attention is memory\-intensive\. On a GPU, moving data between slow off\-chip memory and fast on\-chip memory is expensive\. Naive implementations write the full intermediate score matrix to off\-chip memory, which becomes prohibitive at long sequence lengths\.*Kernel fusion*solves this by combining operations into a single GPU program that keeps intermediate values on\-chip, dramatically reducing memory traffic\.*FlashAttention*\[[1](https://arxiv.org/html/2609.13612#bib.bib1),[2](https://arxiv.org/html/2609.13612#bib.bib2)\]is the canonical fused kernel: a hand\-written CUDA program that underpins nearly every production LLM today\. But it is fixed and inflexible\.

#### The researcher’s dilemma\.

AI researchers constantly invent new attention variants – sliding windows\[[6](https://arxiv.org/html/2609.13612#bib.bib6)\], linear biases\[[7](https://arxiv.org/html/2609.13612#bib.bib7)\], sparse patterns, and more – to improve quality, efficiency, or context length\. Each new variant currently requires expert\-written GPU code to run at useful speeds\. This creates a persistent gap: high\-level Python code is flexible but slow, while production kernels like FlashAttention are fast but inflexible\.

PyTorch 2\.5 introducedflex\_attention\[[14](https://arxiv.org/html/2609.13612#bib.bib14)\], a promising step toward bridging this gap\. Researchers write a small Python function that modifies attention scores, and a compiler turns it into a fused GPU kernel via Triton\[[4](https://arxiv.org/html/2609.13612#bib.bib4),[17](https://arxiv.org/html/2609.13612#bib.bib17)\]\.flex\_attentionis the strongest published bridge between researcher convenience and production performance, supporting sliding windows, ALiBi, and arbitrary score modifications\.

#### A fundamental limitation\.

flex\_attentionhas one important limitation: its modification function operates on the score matrix*after*the multiplicationQ​K⊤QK^\{\\top\}\. Many useful attention variants require transformations*before*this multiplication\. The most consequential example is*Rotary Position Embedding*\(RoPE\)\[[8](https://arxiv.org/html/2609.13612#bib.bib8)\]– the position encoding used by LLaMA, Mistral, Qwen, DeepSeek, and every other recent major LLM\. RoPE rotates queries and keys before multiplication, encoding position information in the geometry of the vectors themselves\. Because the rotation applies toQQandKKseparately, it cannot be expressed as a post\-matmul score modification\. The current workaround applies RoPE in a*separate*kernel before callingflex\_attention, paying two extra memory round\-trips per attention layer \(Figure[1](https://arxiv.org/html/2609.13612#S1.F1)a\)\.

\(a\)flex\_attention:score\_modis post\-matmulQ, K, Vrotate\(Q,K\)Q’K’⊤score\_mod\(s\)softmax\(P\) @ V→\\;\\to\\;output\(b\) AttnFuse:rope\(\)is a pre\-matmul combinatorQ, K, Vrope\(Q,K\)→\\toS=Q’K’⊤→\\tomask→\\tosoftmax→\\to@Vsingle fused Triton kernel

Figure 1:The abstraction gap\. \(a\)flex\_attentionapplies RoPE in a separate kernel becausescore\_modis post\-matmul\. \(b\) AttnFuse represents RoPE as a first\-class pre\-matmul operation and fuses everything into one kernel\.AttnFuseis a domain\-specific language \(DSL\) for attention that makes pre\-multiplication transformations first\-class operations\. Researchers compose ten high\-level building blocks to describe an attention variant declaratively – for example,softmax\(causal\(rope\(Q, K\)\)\) @ V\. AttnFuse’s compiler then emits a single fused GPU kernel for the entire program\. This is a strictly more general abstraction that includes pre\-matmul fusions*by construction*\.

#### Contributions\.

We make three contributions:

- •A compilable DSL for attention\.A ten\-combinator DSL with a two\-level IR and four\-pass compiler guaranteeing single\-kernel fusion for any well\-formed program\.
- •Fused RoPE\.Pre\-matmul fusion of RoPE into attention inner loop, achieving2\.10×2\.10\\timesoverflex\_attentionon causal\+RoPE atN=4096N=4096on RTX 3090\.
- •The Rotation Calculus\.A quantitative principle linking fusion profitability to compute\-to\-bandwidth ratio: fusion wins on Ampere; crossover atN∗≈5,500N^\{\*\}\\approx 5\{,\}500on Hopper, derived algebraically and confirmed experimentally\.

The paper also presents a complete engineering artifact –Flash Decoding, backward pass, block\-sparse attention, HuggingFace integration, and a documented Triton limitation – validating the DSL on real workloads, including Llama\-3\-8B training within 5% of PyTorch’s hand\-tuned backend on H100\.

## 2Background

#### FlashAttention and online softmax\.

FlashAttention\[[1](https://arxiv.org/html/2609.13612#bib.bib1),[2](https://arxiv.org/html/2609.13612#bib.bib2)\]introduced a key algorithmic insight: the softmax normalization\[[5](https://arxiv.org/html/2609.13612#bib.bib5)\]can be computed incrementally, tile by tile, without materializing the fullN×NN\\times Nscore matrix\. The algorithm maintains a running maxmm, a running sumℓ\\ell, and a partial output accumulator across tiles of keys\. When a new tile arrives, the running state is rescaled to account for any change in the maximum\. This recurrence is exact and reduces the number of off\-chip memory accesses from𝒪⁡\(N2\)\\mathcal\{O\}\(N^\{2\}\)to𝒪⁡\(N\)\\mathcal\{O\}\(N\)per query row\. In online softmax, the running state\(m,ℓ,acc\)\(m,\\ell,\\text\{acc\}\)lives entirely in GPU registers, so the full score matrix is never written to off\-chip memory\. This is the foundation of efficient attention implementation\.

#### Why fusion is not always profitable\.

Fusing operations saves off\-chip memory traffic but often increases the amount of computation performed on\-chip\. Whether a fusion is profitable depends on the GPU’s balance between compute throughput and memory bandwidth\. This balance is captured by the*roofline*model: a kernel’s performance is limited by whichever resource is the bottleneck\.

#### Triton\.

Triton\[[4](https://arxiv.org/html/2609.13612#bib.bib4)\]is a language for writing GPU kernels at a tile level, higher\-level than raw CUDA but lower\-level than PyTorch\. It compiles to GPU machine code through an MLIR\-based backend\. AttnFuse generates Triton code; we do not modify the Triton compiler itself, though we identify a lowering limitation in version 3\.3\.1 that creates a useful negative result \(Section[7\.2](https://arxiv.org/html/2609.13612#S7.SS2)\)\.

#### Rotary Position Embeddings \(RoPE\)\.

RoPE\[[8](https://arxiv.org/html/2609.13612#bib.bib8)\]encodes position by rotating each pair of dimensions inQQandKKby an angle proportional to token position, making the dot product⟨Qi′,Kj′⟩\\langle Q^\{\\prime\}\_\{i\},K^\{\\prime\}\_\{j\}\\rangledepend only on relative positionj−ij\-i\. Because the rotation applies toQQandKKseparately*before*Q​K⊤QK^\{\\top\}, RoPE is a pre\-matmul transformation and therefore impossible to express inflex\_attention’s post\-matmulscore\_mod\(Figure[1](https://arxiv.org/html/2609.13612#S1.F1)\)\.

## 3System Design

AttnFuse is designed for AI researchers to implement new attention variants in Python and get a fused GPU kernel that matches production performance\. Figure[2](https://arxiv.org/html/2609.13612#S3.F2)shows the four\-layer pipeline\. The user writes a function using ten composable building blocks; AttnFuse traces the function once per configuration, compiles it through four passes, and dispatches the compiled kernel at runtime with negligible overhead\.

User DSLCodeHigh\-LevelIRTritonKernelRuntimeDispatch

Figure 2:AttnFuse’s four\-layer pipeline: user DSL→\\tohigh\-level IR→\\toTriton kernel→\\toruntime dispatch\.### 3\.1The Researcher Interface

AttnFuse exposes ten combinators that correspond to common attention operations:scaled\_dot\_product,rope,causal,sliding\_window,full,block\_sparse,alibi,additive\_bias,softmax, andrelu\_attention\. A researcher defines an attention variant by composing these combinators:

@af\.attention

defllama\_attn\(Q,K,V\):

s=af\.rope\(Q,K\)

s=af\.causal\(s\)

returnaf\.softmax\(s\)@V

The@af\.attentiondecorator traces this function once, capturing the computation graph\. The DSL is intentionally narrow: it cannot expressflex\_attention’s arbitraryscore\_modcallback\. This is by design – the narrowness is what lets AttnFuse guarantee that every well\-formed program compiles to a single fused kernel\. Adding a new combinator likerope\(\)requires defining its semantics and its compiler effect, but this is a one\-time effort; researchers using the DSL never touch the compiler\.

### 3\.2From DSL to Kernel

The compiler transforms the user’s computation graph through four passes and two intermediate representation \(IR\):

1. 1\.Fuse:Recognizes the canonical attention pattern – score computation, masking, normalization, and value aggregation – and verifies that the graph is fusible\.
2. 2\.Tile:Selects tile sizes \(block dimensions, number of warps, pipeline stages\) appropriate for the target GPU architecture and attention variant\. Tuning tables are precomputed from the sweeps described in Section[5\.2](https://arxiv.org/html/2609.13612#S5.SS2)\.
3. 3\.Lower:Converts the graph to a flatTiledKernelrecord with all constants needed for code generation\.
4. 4\.Codegen:Emits Triton source\. The generated kernel has no runtime branching, all variant\-specific logic is resolved at compile time via Triton’sconstexprmechanism\.

The high\-level IR uses six node types:TensorSymfor input tensors,ScoreOpfor score computation,MaskOpandBiasOpfor modifications,NormOpfor softmax variants, andMatMulPVfor the final value aggregation\. Each graph carries a SHA\-1 signature over its structural content, which serves as a cache key for compiled kernels\. The low\-level IR captures tile shapes, architecture\-specific flags, and constants likeROPE\_KIND,MASK\_KIND, andSAVE\_L\.

### 3\.3Runtime Dispatch

At runtime,run\_attention\(\)computes the graph’s signature, looks up a precompiledLaunchBundle, and invokes the kernel\. Three fast paths bypass the general kernel:

- •Flash Decoding:When the query batch size is 1 \(autoregressive inference\), the KV cache is split across multiple programs for better parallelism \(Section[4\.2](https://arxiv.org/html/2609.13612#S4.SS2)\)\.
- •Block\-sparse:When a sparse block mask is supplied, the kernel iterates only over active blocks \(Section[4\.3](https://arxiv.org/html/2609.13612#S4.SS3)\)\.
- •Hopper spike:On H100 GPUs, a tile\-swept kernel variant is used for causal attention \(Section[5\.2](https://arxiv.org/html/2609.13612#S5.SS2)\)\.

Steady\-state dispatch overhead is under1​μ1\\,\\mus, so the runtime cost is negligible relative to the kernel execution time\.

## 4Key Optimizations

This section describes four optimizations that demonstrate AttnFuse’s expressiveness and performance\.

### 4\.1Fused RoPE

When a researcher includesrope\(Q, K\)in their DSL program, AttnFuse’s code generator produces a kernel that rotates the query matrix once, before the outer loop, and rotates each key tile on the fly inside the inner loop\. The rotation follows the standard NeoX convention: each pair of dimensions is rotated by an angle that depends on the token’s position\. This is implemented using a precomputed index offset and sign multiplier to access the rotated half of each vector\.

### 4\.2Flash Decoding

For autoregressive decoding, the query batch size is 1\. The standard FlashAttention\-2 kernel launches onlyB⋅HqB\\cdot H\_\{q\}programs \(e\.g\., 32 for Llama\-3\-8B\), which underutilizes modern GPUs with hundreds of streaming multiprocessors\.

AttnFuse’s Flash Decoding path addresses this by splitting the KV cache across multiple programs\. Phase 1 launches one program per chunk of the KV cache; each program computes a partial\(m,ℓ,acc\)\(m,\\ell,\\text\{acc\}\)triple using the online\-softmax loop\. Phase 2 runs a small combine kernel that merges these partials via log\-sum\-exp to produce the final output\. For models using Grouped\-Query Attention, we pad the number of query heads per program to 16 by cyclic replication, loadingKKandVVexactly once per program while recovering full tensor\-core throughput\. On Llama\-3\-70B with a 32k cache, this optimization reduces latency from3,153​μ3\{,\}153\\,\\mus to184​μ184\\,\\mus—a17×17\\timesspeedup, beatingflex\_attention’s196​μ196\\,\\mus\.

### 4\.3Block\-Sparse Attention

Many efficient attention variants use sparse masks to reduce quadratic complexity\. AttnFuse supports block\-sparse attention: the user supplies a Python predicate over block coordinates;create\_block\_maskevaluates it once and returns a CSR\-styleBlockMaskwith per\-row active\-block index lists\. Both forward and backward kernels iterate only over active blocks, producing genuinely sub\-quadratic FLOPs\.

On the BigBird mask\[[9](https://arxiv.org/html/2609.13612#bib.bib9)\]atN=4096N\{=\}4096\(∼7\.7%\\sim 7\.7\\%active blocks\), AttnFuse forward runs at84​μ84\\,\\mus vs\.flex\_attention’s270​μ270\\,\\mus, a3\.21×3\.21\\timesspeedup\. The advantage is structural:flex\_attentionpays full𝒪⁡\(N2\)\\mathcal\{O\}\(N^\{2\}\)FLOPs and masks inactive tiles to−∞\-\\infty, while AttnFuse pays proportionally to the active fraction\.

### 4\.4Backward Pass

AttnFuse implements a full backward pass following the FlashAttention\-2 design\. The backward decomposes into three kernels: a preprocessing kernel that computes row\-reduced gradients, a kernel that accumulates gradients with respect to keys and values, and a kernel that accumulates gradients with respect to queries\. This split avoids atomic accumulation: each kernel writes to a clean per\-program output tile\. The forward pass saves a single scalar per query row \(L=m\+log⁡ℓL=m\+\\log\\ell\), which is sufficient to re\-derive softmax probabilities tile\-by\-tile during the backward pass without ever materializing the full score matrix\.

## 5A Cross\-Architecture Analysis

The previous sections established that AttnFuse can express and fuse RoPE\. But is fusion always the right choice? This section shows that the answer depends on the hardware\. We report a structured investigation on H100 NVL \(sm\_90\) that begins with a porting failure, proceeds through systematic tuning, and culminates in a general principle we call the*Rotation Calculus*: whether fused RoPE beats pre\-rotation is determined by the GPU’s compute\-to\-bandwidth ratio\.

### 5\.1The Hopper Gap

Porting the Ampere\-kernel template to H100 NVL produces a2\.10×2\.10\\timesslowdown vs\.flex\_attentionon causal forward atN=4096N\{=\}4096\. We profiled both kernels with Nsight Compute; Table[1](https://arxiv.org/html/2609.13612#S5.T1)reports the counter triangulation\. The production kernel also shows higher register pressure \(217 vs\. 255 regs/thread\) and lower shared memory usage \(65 vs\. 113 KB/block\), consistent with its less aggressive tiling\.

A key reframing emerges from these numbers:flex\_attention’s HMMA pipe sits at32\.6%32\.6\\%on H100, not the6060–70%70\\%that CUTLASS\-class Hopper kernels reach\. The FA\-2 algorithm at this shape is structurally bounded by its non\-matmul fraction \(online softmax, mask logic, address arithmetic\); even fully optimized codegen cannot meaningfully exceed∼35%\\sim 35\\%HMMA\. This reframes the Hopper goal from “WGMMA codegen needed” to “match flex’s ceiling, then look for algorithm\-level gains\.”

Table 1:Nsight Compute counters for causal forward on H100 NVL \(B=4B\{=\}4,H=12H\{=\}12,N=4096N\{=\}4096,D=64D\{=\}64, fp16\)\. The production Ampere\-template kernel reaches15\.8%15\.8\\%HMMA pipe utilization; the sweep\-tuned spike reaches29\.4%29\.4\\%, within3\.23\.2percentage points offlex\_attention’s32\.6%32\.6\\%ceiling\.
### 5\.2Sweep\-Tuning for Hopper

We executed a 16\-configuration sweep overnum\_stages∈\{2,3,4\}\\text\{num\\\_stages\}\\in\\\{2,3,4\\\},BLOCK\_N∈\{64,128,256\}\\text\{BLOCK\\\_N\}\\in\\\{64,128,256\\\}, andnum\_warps∈\{4,8\}\\text\{num\\\_warps\}\\in\\\{4,8\\\}withBLOCK\_M=128\\text\{BLOCK\\\_M\}\{=\}128fixed\. Table[2](https://arxiv.org/html/2609.13612#S5.T2)summarizes the results\.

Table 2:16\-configuration sweep winners on H100 NVL\.Three findings generalize:

1. 1\.The sparse\-table heuristic inverts on Hopper\.On Ampere, causal variants wantBN=32\\text\{BN\}\{=\}32for more programs per SM\. On Hopper, the winner isBN=64\\text\{BN\}\{=\}64for plain causal andBN=128\\text\{BN\}\{=\}128for RoPE\+causal\.
2. 2\.num\_warps=8\\text\{num\\\_warps\}\{=\}8strictly dominates\.num\_warps=4\\text\{num\\\_warps\}\{=\}4is22–3×3\\timesslower in all sweep entries\. Hopper wants one full warp\-group per program\.
3. 3\.num\_stages=3\\text\{num\\\_stages\}\{=\}3vs\. 4 is essentially flat\.Pipeline depth matters less than getting the matmul shape right\.

The sweep\-tuned kernel is packaged as a dispatch\-time fast path \(Table[3](https://arxiv.org/html/2609.13612#S5.T3)\) that activates on H100 for causal attention, closing85%85\\%of the gap toflex\_attentionvia tile selection alone\.

Table 3:H100 ablation – causal forward atB=4B\{=\}4,H=12H\{=\}12,N=4096N\{=\}4096,D=64D\{=\}64, fp16\. Tile selection alone closes85%85\\%of the gap toflex\_attention\.
### 5\.3The Rotation Calculus

We now extend the spike to fused RoPE\+causal and measure the cross\-architecture comparison\. These measurements reveal a general principle: the Rotation Calculus\. Whether to fuse RoPE inside the attention kernel or pre\-rotateQQandKKseparately is determined by the platform’s compute\-to\-bandwidth ratio\. Pre\-rotation pays an𝒪⁡\(N​H​D\)\\mathcal\{O\}\(NHD\)HBM round\-trip once; in\-kernel fusion pays an𝒪⁡\(N2​D/BM\)\\mathcal\{O\}\(N^\{2\}D/\\text\{BM\}\)rotation cost inside the FA\-2 causal loop\. On a bandwidth\-bound platform \(Ampere,≈152\\approx 152FLOPs/byte ridge\), fusion wins decisively\. On a compute\-rich platform \(Hopper,≈300\\approx 300FLOPs/byte ridge\), the crossover occurs betweenN=4​kN\{=\}4\\text\{k\}andN=8​kN\{=\}8\\text\{k\}\. Table[4](https://arxiv.org/html/2609.13612#S5.T4)shows the headline result\.

Table 4:The Rotation Calculus\. Speedup \(×\\times\) of fused RoPE overflex\_attention\+pre\-rotate for causal\+RoPE forward, fp16\. Values\>1\.0\>1\.0indicate AttnFuse fusion wins\. Crossover atN⋆≈5,500N^\{\\star\}\\approx 5\{,\}500\.#### Algebraic identification of the crossover\.

The pre\-rotation cost isTpre=4​B​H​N​D⋅2/BWT\_\{\\text\{pre\}\}=4BHND\\cdot 2/\\text\{BW\}seconds \(readQ,KQ,K, writeQ′,K′Q^\{\\prime\},K^\{\\prime\}\)\. The in\-kernel rotation cost isTfused≈Ntiles⋅crotT\_\{\\text\{fused\}\}\\approx N\_\{\\text\{tiles\}\}\\cdot c\_\{\\text\{rot\}\}whereNtiles=B​H​N2/\(BM⋅BN\)N\_\{\\text\{tiles\}\}=BHN^\{2\}/\(\\text\{BM\}\\cdot\\text\{BN\}\)andcrotc\_\{\\text\{rot\}\}is the per\-tile rotation cost in cycles \(∼800\\sim 800cycles on Hopper\)\. SettingTpre=TfusedT\_\{\\text\{pre\}\}=T\_\{\\text\{fused\}\}yields:

N∗≈8​D⋅BM⋅BN⋅FLOPSBW⋅crot\.N^\{\*\}\\approx\\sqrt\{\\frac\{8D\\cdot\\text\{BM\}\\cdot\\text\{BN\}\\cdot\\text\{FLOPS\}\}\{\\text\{BW\}\\cdot c\_\{\\text\{rot\}\}\}\}\.\(1\)Plugging in H100 NVL numbers\[[16](https://arxiv.org/html/2609.13612#bib.bib16)\]\(BW=3\.9\\text\{BW\}=3\.9TB/s,FLOPS=1015\\text\{FLOPS\}=10^\{15\},BM=128\\text\{BM\}=128,BN=128\\text\{BN\}=128,D=64D=64,crot=800c\_\{\\text\{rot\}\}=800\) givesN∗≈6000N^\{\*\}\\approx 6000, consistent with the measured crossover between 4k and 8k\.

### 5\.4Nsight Compute Profile of the RoPE Spike

To characterize where the in\-kernel rotation cost goes on H100, we profiled the RoPE\+causal spike against the plain\-causal spike \(Table[5](https://arxiv.org/html/2609.13612#S5.T5)\)\. The diagnostic combination is long\-scoreboard rising5\.5×5\.5\\timeswhile DRAM utilization*decreases*: many small HBM\-latency\-blocked loads \(cos, sin,Krot\_halfK\_\{\\text\{rot\\\_half\}\}\) without bandwidth saturation\. This motivated a targeted optimization attempt \(documented in Section[7\.2](https://arxiv.org/html/2609.13612#S7.SS2)\)\.

Table 5:Plain causal vs\. RoPE\+causal forward, H100 NVL,N=4096N\{=\}4096, fp16\.CounterPlainRoPEHMMA pipe %29\.416\.6SM throughput %58\.638\.5Warp occupancy %24\.112\.5Wait stalls % \(math\-pipe\)13\.426\.3Long\-scoreboard % \(HBM\)3\.117\.5DRAM throughput %3\.82\.2Regs/thread120215

## 6Evaluation

All measurements use PyTorch 2\.5\.1, CUDA 12\.1, Triton 3\.1\.0 \(Ampere\) or 3\.3\.1 \(Hopper\)\. Latency is the median of 40–50 CUDA\-event\-timed launches after 8–12 warmups\. Evaluation hardware: RTX 3090 \(sm\_86, 24 GB GDDR6X, 82 SMs, 142 TFLOPS fp16 peak\); H100 NVL \(sm\_90, 132 SMs,∼\\sim1000 TFLOPS fp16 peak via WGMMA\)\. Three baselines:naive\(PyTorch eager\-mode\),sdpa\(torch\.nn\.functional\. scaled\_dot\_product\_attention\), andflex\_attention\(PyTorch 2\.5’s compiled attention\)\.

### 6\.1Forward Head\-to\-Head on Ampere

Table[6](https://arxiv.org/html/2609.13612#S6.T6)reports AttnFuse vs\.flex\_attentionon GPT\-2\-small geometry \(B=4,H=12,D=64B\{=\}4,H\{=\}12,D\{=\}64\) in fp16\. Cells show speedup=flex/AttnFuse=\\text\{flex\}/\\text\{AttnFuse\}\(bold==AttnFuse wins\)\. AttnFuse wins 12 of 16 cells, with sliding\-window attention showing a clean sweep\.

Table 6:Forward latency ratio on RTX 3090, fp16\.
### 6\.2RoPE on Ampere

Becauseflex\_attention’sscore\_modis a post\-matmul hook by design \(Section[4\.1](https://arxiv.org/html/2609.13612#S4.SS1)\), it applies RoPE host\-side in two extra kernels\. Table[7](https://arxiv.org/html/2609.13612#S6.T7)reports the resulting speedup\. AttnFuse wins every cell at1\.651\.65–2\.10×2\.10\\times\.

Table 7:AttnFuse fused RoPE vs\.flex\_attention\+pre\-rotation\. RTX 3090, fp16\.
### 6\.3KV\-Cache Decoding

Table[8](https://arxiv.org/html/2609.13612#S6.T8)reports AttnFuse Flash Decoding vs\.flex\_attentionon three production LLM geometries withQ\.N=1Q\.N\{=\}1\. The Llama\-3\-70B cell shows the17×17\\timesimprovement over the unsplit AttnFuse kernel\.

Table 8:Flash Decoding \(Q\.N=1Q\.N\{=\}1, large KV cache\)\.
### 6\.4H100 Causal Ablation

The H100 spike ablation \(Table[3](https://arxiv.org/html/2609.13612#S5.T3)\) shows the progression from the production Ampere\-tuned kernel \(0\.9310\.931ms\) through the FA\-2 causal split \(0\.6940\.694ms\) to the sweep\-tuned winner \(0\.4880\.488ms\), closing85%85\\%of the gap toflex\_attention’s0\.4430\.443ms reference\. The attempted half\-swap optimization regressed to0\.9520\.952ms and was reverted \(discussed in Section[7](https://arxiv.org/html/2609.13612#S7)\)\.

### 6\.5End\-to\-End: Llama\-3\-8B

To validate AttnFuse on a production workload, we benchmark a single Llama\-3\-8BLlamaDecoderLayer\(random\-init weights,Hq=32H\_\{q\}\{=\}32,Hk​v=8H\_\{kv\}\{=\}8,D=128D\{=\}128\) with the HuggingFace transformers\[[15](https://arxiv.org/html/2609.13612#bib.bib15)\]integration\. AttnFuse registers as anattn\_implementationbackend; switching is a one\-line configuration change\.

Table[9](https://arxiv.org/html/2609.13612#S6.T9)reports latencies on H100 NVL\. Forward\-only, AttnFuse is essentially tied with PyTorch’s SDPA backend atN=2048N\{=\}2048, trailing by only51​μ51\\,\\mus \(\+1\.8%\+1\.8\\%\), and remains within8%8\\%atN=4096N\{=\}4096\. For the full training step \(forward, backward, and optimizer step\), AttnFuse is within5\.2%5\.2\\%of SDPA atN=2048N\{=\}2048and within14\.7%14\.7\\%atN=4096N\{=\}4096\.

Table 9:Llama\-3\-8B on H100 NVL, random initialization,B=1B\{=\}1, fp16\. Latencies in milliseconds\.On RTX 3090, AttnFuse achieves1\.06×1\.06\\timesSDPA on the training step;flex\_attentionOOMs at HEAD\_DIM=128, so AttnFuse is the only Triton\-based attention compiler that runs Llama\-3\-8B training on consumer Ampere\. These results demonstrate that a DSL\-emitted kernel can match hand\-tuned CUDA on production workloads\.

#### Two notes on the H100 numbers\.

First, HuggingFace’sLlamaDecoderLayerapplies RoPE toQQandKK*before*calling the registered attention function, so the graph traced is plain causal, not fused RoPE\. Threadingcos\\cos,sin\\sinthrough the HF hook would invoke the fused\-RoPE path, which per the Rotation Calculus should improve the forward number atN≤4​kN\\leq 4\\text\{k\}\. Second, the spike’s dispatch predicate rejects calls withsave\_lseset, so the training\-step forward uses the production kernel; widening predicate coverage would close the training\-step gap\. Both are well\-scoped follow\-on items\.

### 6\.6Correctness

87 GPU pytest cells cover the variant matrix×\\timesdtype×\\timesshape×\\timesGQA configuration\. 80 Hypothesis\-driven\[[11](https://arxiv.org/html/2609.13612#bib.bib11)\]property fuzzer examples randomize configurations across forward and backward\. All gradients are verified against a naive PyTorch reference within FA2’s documented fp16 tolerance \(\|e​r​r\|<2×10−2\|err\|<2\\times 10^\{\-2\}\)\.

## 7Discussion

The results reframe three common assumptions about attention codegen for modern GPUs and provide practical guidance for researchers and practitioners\.

#### Hopper does not necessarily require WGMMA intrinsics for attention\.

Section[5](https://arxiv.org/html/2609.13612#S5)measuredflex\_attention’s HMMA pipe ceiling on H100 at32\.6%32\.6\\%, not the6060–70%70\\%that CUTLASS\-class Hopper kernels sometimes reach\. The FA\-2 non\-matmul fraction – online softmax, mask logic, address arithmetic – is the binding constraint at this shape, not the WGMMA\-vs\-HMMA codegen choice\. Our sweep\-tuned spike reaches29\.4%29\.4\\%HMMA via tile selection alone, within3\.23\.2percentage points of this ceiling\. Closing the residual gap requires FA\-3\-style algorithmic restructure \(producer/consumer warp specialization that overlaps softmax with the next matmul\), not explicit Hopper intrinsics\. For AI systems researchers, this means that significant Hopper gains for attention may come from algorithm redesign rather than low\-level tuning alone\.

#### Fusion is platform\-dependent, and that’s quantifiable\.

The Rotation Calculus is the explicit form of an intuition the FlashAttention community has long gestured at: as accelerators rebalance toward compute, the value of HBM\-saving fusions shrinks\. We measure and quantify the crossover for fused RoPE; the same algebra applies to any pre\-matmul fusion whose in\-kernel cost grows withNNfaster than its pre\-pass analogue\. For practitioners, this provides a concrete rule: on compute\-rich hardware like H100, pre\-rotation may be preferable for long contexts; on bandwidth\-bound hardware like Ampere, fusion is almost always the right choice\.

#### A high\-level DSL can match hand\-tuned CUDA on production LLMs\.

The51​μ51\\,\\mus gap from SDPA atN=2048N\{=\}2048forward on H100, and the5%5\\%training\-step gap at the same shape, are the production\-grade claim\. The DSL surface is what makes adding the next variant a one\-day change rather than a multi\-week kernel rewrite\. This suggests that the historical tradeoff between researcher productivity and production performance is not fundamental; attention\-specific compilation can offer both\.

### 7\.1Workload Positioning

Where is AttnFuse the right choice today? Based on our results, AttnFuse is the*strong choice*for LLM fine\-tuning at22k–44k tokens \(within5%5\\%of SDPA on training,2\.10×2\.10\\timesoverflex\_attentionon RTX 3090 for RoPE\), KV\-cache inference at44k–3232k cache lengths \(17×17\\timesspeedup over unsplit AttnFuse, parity withflex\_attentionon Llama\-3\-70B at 32k\), and block\-sparse training \(3\.21×3\.21\\timesoverflex\_attentionon BigBird atN=4N\{=\}4k\)\. It is*comparable*toflex\_attentionon ALiBi, sliding window, and custom score modifications \(55–10%10\\%range on Ampere\), where the DSL surface is the primary value\. Long\-context RoPE on Hopper \(N≥16N\\geq 16k\) is*future work*: the Rotation Calculus favors pre\-rotation at these lengths, and an FA\-3\-style producer/consumer kernel is the path forward\. For full pre\-training at44k–88k, the picture is*mixed*: forward parity on H100, but the backward uses the production kernel \(spike predicate rejectssave\_lse\-set calls\), and RoPE crosses over nearN≈6N\\approx 6k\.

### 7\.2Lessons Learned: A Negative Result

The investigation of fused RoPE on Hopper revealed a compiler limitation worth documenting\. The NCU diagnosis of the RoPE spike \(Section[5](https://arxiv.org/html/2609.13612#S5)\) pointed at the per\-tileKrot\_halfK\_\{\\text\{rot\\\_half\}\}HBM load as the dominant new stall: long\-scoreboard stalls rose5\.5×5\.5\\timesfrom3\.1%3\.1\\%to17\.5%17\.5\\%, while DRAM utilization actually decreased\. TheKrot\_halfK\_\{\\text\{rot\\\_half\}\}tile isKKwith the twoDD\-halves swapped; algebraically, it can be derived fromKKby a register layout permutation with no second HBM access\.

We implemented this via Triton’stl\.reshape\+tl\.permute\+tl\.split\+tl\.joinchain\. Table[10](https://arxiv.org/html/2609.13612#S7.T10)shows the bisected result\. Prong B achieved the intended reduction in HBM dependency: long\-scoreboard stalls fell from17\.45%17\.45\\%to3\.33%3\.33\\%\. However, it introduced a stronger SMEM dependency: short\-scoreboard stalls rose from7\.04%7\.04\\%to15\.68%15\.68\\%\. HMMA utilization dropped from16\.6%16\.6\\%to10\.3%10\.3\\%, SM throughput fell from38\.5%38\.5\\%to26\.8%26\.8\\%, and wall\-clock latency regressed by20%20\\%to0\.9520\.952ms, worse than theflex\_attention\+pre\-rotate baseline at0\.8340\.834ms\.

Table 10:Negative result: attempted half\-swap optimization\. Prong B reduced HBM stalls but introduced stronger SMEM stalls, causing a20%20\\%latency regression\.#### Cause and implication\.

The root cause is that Triton 3\.3\.1 lowers thetl\.reshape\+tl\.permute\+tl\.split\+tl\.joinchain on sm\_90 to a SMEM round\-trip, not to a register\-level warp shuffle\. The kernel writes the tile to shared memory, performs the layout permutation there, and reads it back to registers\. The SMEM dependency cost exceeds the HBM dependency cost it replaced because Hopper’s L2\-cached HBM access for theKrot\_halfK\_\{\\text\{rot\\\_half\}\}load is cheaper than the SMEM staging Triton inserts\.

The optimization is right in principle and will pay off as soon as Triton’s lowering recognizes the half\-swap pattern as a warp\-shuffle candidate\. Inline PTX \(\_\_shfl\_xor\_sync\) would also achieve the goal but forfeits Triton’s portability\. The bisected counter delta in Table[10](https://arxiv.org/html/2609.13612#S7.T10)provides a falsifiable test case that the Triton compiler community can address\. For researchers using Triton, this serves as a cautionary tale: high\-level tensor operations on sm\_90 may lower to surprising SMEM round\-trips that defeat the intended optimization\.

## 8Related Work

#### FlashAttention\.

Dao et al\.\[[1](https://arxiv.org/html/2609.13612#bib.bib1),[2](https://arxiv.org/html/2609.13612#bib.bib2)\]pioneered IO\-aware tiled attention, establishing the online softmax recurrence and tiled execution that underpin nearly every efficient attention implementation today\. AttnFuse adopts the same algorithmic foundation but exposes it through a composable DSL rather than a fixed CUDA binary\. FlashAttention\-3\[[3](https://arxiv.org/html/2609.13612#bib.bib3)\]is the Hopper\-targeted hand\-written follow\-on; its producer/consumer warp specialization overlaps softmax with the next matmul, recovering some of the non\-matmul overhead that limits FA\-2 at∼30%\\sim 30\\%HMMA on H100 \(Section[5](https://arxiv.org/html/2609.13612#S5)\)\. Porting AttnFuse’s codegen to the FA\-3\[[3](https://arxiv.org/html/2609.13612#bib.bib3)\]inner\-loop structure is a natural next step for closing the residual Hopper gap\.

#### PyTorchflex\_attention\.

flex\_attention\(PyTorch 2\.5\+\) is the closest published abstraction to AttnFuse\. As discussed in Section[1](https://arxiv.org/html/2609.13612#S1)and Figure[1](https://arxiv.org/html/2609.13612#S1.F1), thescore\_modhook is post\-matmul*by design*; AttnFuse’s combinator surface includes pre\-matmul nodes \(rope\(\)\) by design\. This structural difference is the core distinction:flex\_attentionis a general\-purpose abstraction for modifying scores; AttnFuse is an attention\-specific DSL that guarantees single\-kernel fusion for any well\-formed program, including those with pre\-matmul transformations\.

#### Triton and TorchInductor\.

Triton\[[4](https://arxiv.org/html/2609.13612#bib.bib4)\]provides the JIT compilation infrastructure that both AttnFuse and PyTorch’s TorchInductor build upon\. TorchInductor is PyTorch’s Triton\-emitting backend forflex\_attentionand other compiled operations\. On Hopper, our investigation \(Section[5](https://arxiv.org/html/2609.13612#S5)\) finds that TorchInductor’s WGMMA codegen achieves an HMMA pipe ceiling of32\.6%32\.6\\%at the shape tested\. The Rotation Calculus reframes this as a structural property of the FA\-2 algorithm rather than a Triton limitation—a finding that has implications for the broader Triton ecosystem\. Our documented negative result on Triton 3\.3\.1’s lowering of reshape/permute chains \(Section[7](https://arxiv.org/html/2609.13612#S7)\) offers a falsifiable test case for the Triton compiler community\.

#### CUTLASS and vendor libraries\.

NVIDIA CUTLASS\[[12](https://arxiv.org/html/2609.13612#bib.bib12)\]is a C\+\+ template library for GPU matrix multiplication and attention\. PyTorch’s SDPA backend on H100 reaches CUTLASS\-class throughput\. The matching we report in Section[6\.5](https://arxiv.org/html/2609.13612#S6.SS5)\(Llama\-3\-8B forward within51​μ51\\,\\mus of SDPA atN=2048N\{=\}2048\) demonstrates how close a high\-level DSL can come to hand\-tuned vendor libraries—without requiring the developer to write C\+\+ or reason about CUDA specifics\.

#### Domain\-specific kernel libraries\.

xFormers\[[13](https://arxiv.org/html/2609.13612#bib.bib13)\]provides a library of memory\-efficient attention kernels but is not a compiler—users choose among hand\-written variants\. Other libraries such as FlashAttention’s official repository and various open\-source implementations offer fixed kernels for common patterns\. AttnFuse’s domain\-specificity \(attention only\) is what permits the guaranteed\-fusion compiler contract that a general\-purpose kernel DSL cannot promise\. This narrowness is a feature, not a limitation: it enables the compiler to make strong guarantees about memory behavior and kernel structure\.

## 9Conclusion and Future Work

We presented AttnFuse, an embedded Python DSL for attention that compiles to fused Triton kernels\. The DSL makes RoPE a first\-class pre\-matmul combinator, offering a strictly more general abstraction thanflex\_attention’s post\-matmulscore\_mod\. On RTX 3090, fused RoPE achieves2\.10×2\.10\\timesoverflex\_attention\. On H100, a tile\-swept dispatch path closes85%85\\%of the gap toflex\_attention, with tensor\-core utilization within3\.23\.2percentage points of the structural ceiling\. The Rotation Calculus provides the first systematic quantification of when in\-kernel RoPE fusion is profitable as a function of compute\-to\-bandwidth ratio\. End\-to\-end, AttnFuse runs a Llama\-3\-8B training step within5%5\\%of hand\-tuned SDPA on both architectures\.

### 9\.1Limitations and Future Directions

Several limitations point to natural next steps\. The Hopper spike currently covers causal MHA/GQA atD∈\{64,128\}D\\in\\\{64,128\\\}, fp16/bf16,N≥2048N\\geq 2048, forward\-only\. Extensions include threadingcos,sin\\cos,\\sinthrough the HuggingFace hook for fused RoPE, widening the spike predicate to LSE\-saving forward calls, and extending codegen to sliding\-window, ALiBi, and additive\-bias variants\.

On H100 atN≥8​kN\\geq 8\\text\{k\}, pre\-rotation wins\. Two recovery paths are promising: an FA\-3 producer/consumer kernel that overlaps rotation cost with the next matmul, or Triton compiler support for warp\-shuffle\-lowered half\-swap\. The negative result we documented provides a concrete test case for the Triton compiler community\.

The two\-level IR is structurally ready for Ring Attention\[[10](https://arxiv.org/html/2609.13612#bib.bib10)\]; lowering to a sequence\-parallel kernel would enable scaling across multiple GPUs\. Hopper’s fp8 tensor cores offer a natural extension for low\-precision training and inference; adding fp8 support introduces one new IR concept\.

We believe the Rotation Calculus will outlast the AttnFuse implementation\. As accelerators continue to rebalance toward compute, this calculus will become increasingly relevant for AI systems researchers\.

## References

- \[1\]Tri Dao, Daniel Y\. Fu, Stefano Ermon, Atri Rudra, and Christopher Ré\.FlashAttention: Fast and Memory\-Efficient Exact Attention with IO\-Awareness\.In*Advances in Neural Information Processing Systems \(NeurIPS\)*, 2022\.
- \[2\]Tri Dao\.FlashAttention\-2: Faster Attention with Better Parallelism and Work Partitioning\.*arXiv preprint*, arXiv:2307\.08691, 2023\.
- \[3\]Jay Shah, Ganesh Bikshandi, Ying Zhang, Vijay Thakkar, Pradeep Ramani, and Tri Dao\.FlashAttention\-3: Fast and Accurate Attention with Asynchrony and Low\-Precision\.*arXiv preprint*, arXiv:2407\.08608, 2024\.
- \[4\]Philippe Tillet, H\. T\. Kung, and David Cox\.Triton: An Intermediate Language and Compiler for Tiled Neural Network Computations\.In*Proceedings of the 3rd ACM SIGPLAN International Workshop on Machine Learning and Programming Languages \(MAPL\)*, 2019\.
- \[5\]Maxim Milakov and Natalia Gimelshein\.Online Normalizer Calculation for Softmax\.*arXiv preprint*, arXiv:1805\.02867, 2018\.
- \[6\]Rewon Child, Scott Gray, Alec Radford, and Ilya Sutskever\.Generating Long Sequences with Sparse Transformers\.*arXiv preprint*, arXiv:1904\.10509, 2019\.
- \[7\]Ofir Press, Noah A\. Smith, and Mike Lewis\.Train Short, Test Long: Attention with Linear Biases Enables Input Length Extrapolation\.In*International Conference on Learning Representations \(ICLR\)*, 2022\.
- \[8\]Jianlin Su, Yu Lu, Shengfeng Pan, Ahmed Murtadha, Bo Wen, and Yunfeng Liu\.RoFormer: Enhanced Transformer with Rotary Position Embedding\.*arXiv preprint*, arXiv:2104\.09864, 2021\.
- \[9\]Manzil Zaheer, Guru Guruganesh, Avinava Dubey, Joshua Ainslie, Chris Alberti, Santiago Ontañón, Philip Pham, Anirudh Ravula, Qifan Wang, Li Yang, and Amr Ahmed\.Big Bird: Transformers for Longer Sequences\.In*Advances in Neural Information Processing Systems \(NeurIPS\)*, 2020\.
- \[10\]Hao Liu and Pieter Abbeel\.Ring Attention with Blockwise Transformers for Near\-Infinite Context\.*arXiv preprint*, arXiv:2310\.01889, 2023\.
- \[11\]David R\. MacIver, Zac Hatfield\-Dodds, and many others\.Hypothesis: A New Approach to Property\-Based Testing\.*Journal of Open Source Software*, 4\(43\), 2019\.
- \[12\]Vijay Thakkar, Pradeep Ramani, Cris Cecka, and others\.CUTLASS: CUDA Templates for Linear Algebra Subroutines\.NVIDIA Corporation, 2023\.
- \[13\]Benjamin Lefaudeux, Francisco Massa, Diane Liskovitch, Wenhan Xiong, Vittorio Caggiano, Sean Naren, Min Xu, Jieru Hu, Marta Tintore, Susan Zhang, Patrick Labatut, Daniel Haziza, Luca Wehrstedt, Jeremy Reizenstein, and Grigory Sizov\.xFormers: A Modular and Hackable Transformer Modelling Library\.GitHub repository, 2022\.
- \[14\]PyTorch Contributors\.PyTorch 2\.5 flex\_attention Documentation\.https://pytorch\.org/docs/stable/generated/torch\.nn\.functional\.scaled˙dot˙product˙attention\.html, 2024\.Accessed: 2026\-01\-15\.
- \[15\]Thomas Wolf, Lysandre Debut, Victor Sanh, Julien Chaumond, Clement Delangue, Anthony Moi, Pierric Cistac, Tim Rault, Rémi Louf, Morgan Funtowicz, Joe Davison, Sam Shleifer, Patrick von Platen, Clara Ma, Yacine Jernite, Julien Plu, Canwen Xu, Teven Le Scao, Sylvain Gugger, Mariama Drame, Quentin Lhoest, and Alexander M\. Rush\.HuggingFace Transformers: State\-of\-the\-Art Natural Language Processing\.https://github\.com/huggingface/transformers, 2024\.Accessed: 2026\-01\-15\.
- \[16\]NVIDIA\.NVIDIA H100 Tensor Core GPU Architecture Overview\.https://www\.nvidia\.com/en\-us/data\-center/h100/, 2024\.Accessed: 2026\-01\-15\.
- \[17\]Philippe Tillet and OpenAI\.Triton Compiler and MLIR Backend\.https://github\.com/openai/triton, 2024\.Accessed: 2026\-01\-15\.