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

Varun Kumar Dasoju; Tian Zhao

cs.LGSep 11, 2026

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

Authors: Varun Kumar Dasoju, Tian Zhao

Organizations: Department of Computer Science University of Wisconsin–Milwaukee

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 $\times$ 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

Panel of Figure 1 — 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.

Panel of Figure 2 — Figure 2: AttnFuse’s four-layer pipeline: user DSL $\to$ high-level IR $\to$ Triton kernel $\to$ 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 $\times$
+ FA-2 causal split + bigger tiles	0.694 ms	1.34 $\times$
+ BLOCK_N=64 sweep winner	0.488 ms	1.91 $\times$
Hopper spike (dispatched)	0.488 ms	1.91 $\times$
flex_attention reference	0.443 ms	2.10 $\times$

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 $\backslash$ 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 (

\times

) of fused RoPE over flex_attention +pre-rotate for causal+RoPE forward, fp16. Values

>1.0

indicate AttnFuse fusion wins. Crossover at

N^{\star}\approx 5{,}500

Counter	Plain	RoPE
HMMA pipe %	29.4	16.6
SM throughput %	58.6	38.5
Warp occupancy %	24.1	12.5
Wait stalls % (math-pipe)	13.4	26.3
Long-scoreboard % (HBM)	3.1	17.5
DRAM throughput %	3.8	2.2

Table 5: Plain causal vs. RoPE+causal forward, H100 NVL,

N{=}4096

, fp16.

Variant	N=512	N=1024	N=2048	N=4096
Dense	1.74 $\times$	1.00 $\times$	0.98 $\times$	0.95 $\times$
Causal	1.10 $\times$	1.15 $\times$	1.04 $\times$	0.94 $\times$
SW (W=256)	1.16 $\times$	1.16 $\times$	1.10 $\times$	1.05 $\times$
Causal+ALiBi	1.20 $\times$	1.19 $\times$	1.10 $\times$	0.96 $\times$

Table 6: Forward latency ratio on RTX 3090, fp16.

Variant	N=512	N=1024	N=2048	N=4096
RoPE + Dense	1.99 $\times$	2.05 $\times$	1.92 $\times$	1.78 $\times$
RoPE + Causal	2.10 $\times$	2.04 $\times$	1.95 $\times$	2.10 $\times$
RoPE + SW	1.81 $\times$	1.66 $\times$	1.74 $\times$	1.65 $\times$

Table 7: AttnFuse fused RoPE vs. flex_attention +pre-rotation. RTX 3090, fp16.

Geometry	Cache	AttnFuse	flex
Llama-3-8B (32q/8kv/128)	4096	0.034 ms	0.039 ms
Llama-3-70B (64q/8kv/128)	32768	0.184 ms	0.196 ms
Mistral-7B (32q/8kv/128)	16384	0.072 ms	0.079 ms

Table 8: Flash Decoding (

Q.N{=}1

, large KV cache).

Workload	Sequence	SDPA	flex	AttnFuse
Forward only	$N{=}2048$	2.82	2.71	2.87
Forward only	$N{=}4096$	5.90	5.93	6.36
Forward+backward+step	$N{=}2048$	7.27	7.49	7.65
Forward+backward+step	$N{=}4096$	14.39	16.02	16.51

Table 9: Llama-3-8B on H100 NVL, random initialization,

B{=}1

, fp16. Latencies in milliseconds.

Metric	S7-A Spike	Prong B (failed)	Change
Long-SB stalls (%)	17.45	3.33	-14.12 pp
Short-SB stalls (%)	7.04	15.68	+8.64 pp
HMMA pipe (%)	16.6	10.3	-6.3 pp
SM throughput (%)	38.5	26.8	-11.7 pp
Latency (ms)	0.793	0.952	$+20\%$

Table 10: Negative result: attempted half-swap optimization. Prong B reduced HBM stalls but introduced stronger SMEM stalls, causing a

20\%

latency regression.

Efficient Attention

Efficient attention reduces the compute, memory traffic, and cache costs of Transformer attention while preserving useful long-range interactions.

See all 134 papers

Apr 16, 2026

Nautilus: An Auto-Scheduling Tensor Compiler for Efficient Tiled GPU Kernels

We present Nautilus, a novel tensor compiler that moves toward fully automated math-to-kernel optimization. Nautilus compiles a high-level algebraic specification of tensor operators into efficient tiled GPU kernels. Nautilus's successive lowering design allows high-level optimizations, expression rewrites, and tile optimizations to be jointly applied in a single end-to-end system. Nautilus presents a novel auto-scheduler that discovers sequences of high-level optimizations, while preserving the regular program structure needed by tile optimizers. Nautilus's auto-scheduler captures complex interactions and trade-offs in the high-level optimizations, including aggressive global transformations like advanced reduction fusion. Nautilus is the first end-to-end tensor compiler capable of starting from a math-like description of attention and automatically discovering FlashAttention-3-like kernels, offloading the entire burden of optimization from the programmer to the compiler. Across five transformer-based models and 150 evaluation configurations on NVIDIA GH200 and RTX 5090 GPUs, Nautilus achieves up to 23% higher throughput than state-of-the-art compilers on GH200 and up to 42% on RTX 5090, while matching or exceeding manually written cuDNN kernels on many long-sequence configurations.

Yifan Zhao, Yuchen Yang, Matei Budiu +1

University of Illinois Urbana-Champaign USA

GPU Kernel Optimization GPU Acceleration High-Performance Computing GPU Kernel Generation

Ask PDF

May 6, 2026

CuBridge: An LLM-Based Framework for Understanding and Reconstructing High-Performance Attention Kernels

Efficient CUDA implementations of attention mechanisms are critical to modern deep learning systems, yet supporting diverse and evolving attention variants remains challenging. Existing frameworks and compilers trade performance for flexibility, while expert-written kernels achieve high efficiency but are difficult to adapt. Recent work explores large language models (LLMs) for GPU kernel generation, but prior studies report unstable correctness and significant performance gaps for complex operators such as attention. We present CuBridge, an LLM-based framework that adapts expert-written attention kernels through a structured lift-transfer-lower workflow. CuBridge starts from expert-written CUDA attention kernels and lifts them into an executable intermediate representation that makes execution orchestration explicit while abstracting low-level CUDA syntax. Given a user-provided PyTorch specification, CuBridge generates and verifies a target IR program, then reconstructs optimized CUDA code via reference-guided lowering. Across diverse attention variants and GPU platforms, CuBridge consistently produces correct kernels and substantially outperforms general frameworks, compiler-based approaches, and prior LLM-based methods.

Xing Ma, Yangjie Zhou, Wu Sun +6

Shanghai Jiao Tong University · National University of Singapore · Shanghai Qi Zhi Institute

GPU Kernel Optimization GPU Acceleration Self-Attention Large Language Model-Guided Optimization LLM-Based Program Synthesis Transformer Attention+1

Ask PDF

Mar 2, 2026

TiledAttention: a CUDA Tile SDPA Kernel for PyTorch

TiledAttention is a scaled dot-product attention (SDPA) forward operator for SDPA research on NVIDIA GPUs. Implemented in cuTile Python (TileIR) and exposed as a PyTorch-callable function, it is easier to modify than low-level CUDA templates while retaining realistic behavior via online softmax and tiled

K,V

streaming. Algorithmically, TiledAttention follows the established FlashAttention-style online-softmax formulation; our novelty is the cuTile/TileIR implementation strategy, schedule-level modifiability, and reproducible benchmarking/profiling workflow. The approach is both performant and directly editable at the schedule level from Python (tile shapes, staging, shared-memory layout), enabling rapid, reproducible kernel research without template-heavy CUDA/CUTLASS rewrites. We benchmark TiledAttention on an NVIDIA DGX GB10 node with a reproducible harness and compare against PyTorch SDPA (auto-dispatch), explicit unfused baselines (torch_sdpa_math, standard eager attention), and forced backend probes (FlashAttention2, EffecientAttention, CuDNN Attention) across sequence length, head dimension, and precision (FP16/BF16). While production fused baselines remain stronger overall, TiledAttention delivers large speedups over standard eager attention paths and is available for direct use within PyTorch workflows, providing a practical balance between performance and customizability.

Taimur Khan

Helmholtz Centre for Environmental Research - UFZ

GPU Kernel Optimization Self-Attention Efficient Attention Softmax Attention

Ask PDF

Abstract

Efficient Attention

Explore similar work

Nautilus: An Auto-Scheduling Tensor Compiler for Efficient Tiled GPU Kernels

CuBridge: An LLM-Based Framework for Understanding and Reconstructing High-Performance Attention Kernels

TiledAttention: a CUDA Tile SDPA Kernel for PyTorch