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× speedup 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.
Figures & tables
Figure 1: The abstraction gap. (a) flex_attention applies RoPE in a separate kernel because score_mod is post-matmul. (b) AttnFuse represents RoPE as a first-class pre-matmul operation and fuses everything into one kernel.
Figure 2: AttnFuse’s four-layer pipeline: user DSL → high-level IR → Triton kernel → runtime dispatch.
Metric
Production
Spike
flex_attention
HMMA pipe (%)
15.8
29.4
32.6
SM throughput (%)
35.5
58.6
43.7
Warp occupancy (%)
12.0
24.1
12.0
DRAM throughput (%)
2.1
3.8
4.3
Wait stalls (%)
29.9
13.4
21.5
Long-SB stalls (%)
1.2
3.1
—
Table 1: Nsight Compute counters for causal forward on H100 NVL ( B=4 , H=12 , N=4096 , D=64 , fp16). The production Ampere-template kernel reaches 15.8% HMMA pipe utilization; the sweep-tuned spike reaches 29.4% , within 3.2 percentage points of flex_attention ’s 32.6% ceiling.
Variant
Winner
Latency
Plain causal
BN=64, nw=8, ns=3
0.488 ms
RoPE+causal
BN=128, nw=8, ns=3
0.779 ms
Table 2: 16-configuration sweep winners on H100 NVL.
Configuration
Latency
vs. prod.
Production Ampere kernel on H100
0.931 ms
1.00 ×
+ FA-2 causal split + bigger tiles
0.694 ms
1.34 ×
+ BLOCK_N=64 sweep winner
0.488 ms
1.91 ×
Hopper spike (dispatched)
0.488 ms
1.91 ×
flex_attention reference
0.443 ms
2.10 ×
Table 3: H100 ablation – causal forward at B=4 , H=12 , N=4096 , D=64 , fp16. Tile selection alone closes 85% of the gap to flex_attention .
GPU \ N
512
1024
2048
4096
8192
16384
RTX 3090
2.10
2.04
1.95
2.10
—
—
H100 NVL
—
—
1.37
1.05
0.82
0.69
Table 4: The Rotation Calculus. Speedup ( × ) of fused RoPE over flex_attention +pre-rotate for causal+RoPE forward, fp16. Values >1.0 indicate AttnFuse fusion wins. Crossover at N⋆≈5,500 .