TiledAttention: a CUDA Tile SDPA Kernel for PyTorch

Taimur Khan

cs.LGMar 2, 2026

TiledAttention: a CUDA Tile SDPA Kernel for PyTorch

Authors: Taimur Khan

Organizations: Helmholtz Centre for Environmental Research - UFZ

Abstract

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.

Figures & tables

Load

Q

tile

\rightarrow

Stream

K,V

tiles + score/mask

\rightarrow

Update

(m,\ell,o)

; normalize + store

Figure 1: Figure 2: TiledAttention forward pipeline at a glance.

Parameter	Typical values	Effect
$T_{M}$ (Q rows)	64, 128	CTA work per launch; affects occupancy and reuse
$T_{N}$ (KV columns)	64, 128, 256	Streaming granularity; impacts memory coalescing and softmax update cost
Stages / prefetch	1–4	Hides global memory latency via pipelining
Shared-mem layout	row/col swizzles	Impacts bank conflicts and vectorized loads

Table 1: Tuning/search parameters exposed by TiledAttention (representative).

Item	Value
System	NVIDIA DGX GB10 host GPU
GPU driver	580.126.09
CUDA toolkit	CUDA 13.1 (V13.1.115)
cuTile	1.1.0
PyTorch	2.10.0+cu130
Timing method	CUDA events; median + p95 over $N_{r}$ runs

Table 2: Reproducibility checklist for the measured runs in this artifact.

Comparison	Mean ratio	Median ratio	Win count
TiledAttention / PyTorch SDPA (auto)	0.632 $\times$	0.634 $\times$	4 / 80
TiledAttention / PyTorch SDPA (math)	28.15 $\times$	20.31 $\times$	80 / 80
TiledAttention / standard eager attention	14.36 $\times$	10.76 $\times$	80 / 80

Table 3: Aggregate throughput summary over 80 study points (tokens/s ratio).

Regime	Shapes	Best $(T_{M},T_{N})$	Runner-up	Sensitivity
Short $S$	$S{=}1024,D{=}128$	(64, 128)	(64, 64)	3.92%
Mid $S$	$S{=}4096,D{=}128$	(64, 128)	(64, 64)	8.53%
Long $S$	$S{=}8192,D{=}128$	(64, 128)	(64, 64)	9.39%

Table 4: Best tile settings by regime from the measured tuning sweep.

Config	$(1024,64)$	$(2048,64)$	$(4096,128)$
async_default (fp32)	0.766 $\times$	0.668 $\times$	1.017 $\times$
tm64_tn64 (fp16acc)	0.779 $\times$	0.698 $\times$	0.968 $\times$
async_auto (shape-aware)	0.775 $\times$	0.694 $\times$	1.017 $\times$

Table 5: Reduced benchmark ({(1024,64), (2048,64), (4096,128)}, FP16, non-causal): ratio to PyTorch SDPA (auto-dispatch) throughput.

Priority	Recommended path	Rationale
Peak production throughput across broad shapes	PyTorch SDPA (auto)	Best aggregate performance in this study.
Rapid kernel iteration / schedule research	TiledAttention	Python-level tile policy changes with reproducible profiling loop.
Custom kernel development in a broader DSL ecosystem	Triton / FlashAttention codebases	Strong alternatives when teams can afford separate kernel engineering and tuning effort.
Custom masking/layout experiments	TiledAttention	Easier to modify kernel behavior than template-heavy CUDA paths.
Replacing unfused attention baselines	TiledAttention	Large gains vs eager and math- SDPA paths (Table 3 ).

Table 6: Decision guide for HPC deployment.

GPU Kernel Optimization

GPU kernel optimization improves the speed and memory efficiency of computations by tailoring code to GPU hardware and workloads.

See all 202 papers

Evaluating CUDA Tile for AI Workloads on Hopper and Blackwell GPUs
Apr 25, 2026Divakar Kumar Yadav, Tian Zhao, Deepak KumarEfficient Transformer Inference GPU Kernel Optimization
University of Wisconsin–Milwaukee Milwaukee, WI, USA · Illinois Institute of Technology Chicago, IL, USA
AttnFuse: A Composable DSL for Compiling Attentions to Fused GPU Kernels
Sep 11, 2026Varun Kumar Dasoju, Tian ZhaoEfficient Attention GPU Kernel Optimization
Department of Computer Science University of Wisconsin–Milwaukee
Nautilus: An Auto-Scheduling Tensor Compiler for Efficient Tiled GPU Kernels
Apr 16, 2026Yifan Zhao, Yuchen Yang, Matei Budiu +1GPU Kernel Optimization GPU Acceleration
University of Illinois Urbana-Champaign USA

Apr 25, 2026J/K move · Enter open · S save

Evaluating CUDA Tile for AI Workloads on Hopper and Blackwell GPUs

Divakar Kumar Yadav, Tian Zhao, Deepak Kumar

University of Wisconsin–Milwaukee Milwaukee, WI, USA · Illinois Institute of Technology Chicago, IL, USA

Efficient Transformer Inference GPU Kernel Optimization GPU Acceleration

Open paper AskPDF arXiv

NVIDIA's CUDA Tile (CuTile) introduces a Python-based, tile-centric abstraction for GPU kernel development that aims to simplify programming while retaining Tensor Core and Tensor Memory Accelerator (TMA) efficiency on modern GPUs. We present the first independent, cross-architecture evaluation of CuTile against established approaches such as cuBLAS, Triton, WMMA, and raw SIMT on three NVIDIA GPUs spanning Hopper and Blackwell: H100 NVL, B200, and RTX PRO 6000 Blackwell Server Edition. We benchmark representative AI workloads, including GEMM, fused multi-head attention, and end-to-end LLM inference in BF16/FP16 precision, to assess both performance and portability. Our results show that CuTile effectiveness is strongly workload- and architecture-dependent. On datacenter-class Blackwell (B200), CuTile achieves up to 1007 TFLOP/s for fused attention, outperforming FlashAttention-2 by 2.5x while requiring only 60 lines of Python kernel code. For GEMM, CuTile reaches 52-79% of cuBLAS performance in 22 lines of code (versus 123 for WMMA), making it a practical replacement for hand-written CUDA kernels but not yet for vendor-optimized libraries. However, the same CuTile attention kernel achieves only 53% of FlashAttention-2 throughput on RTX PRO 6000 (sm_120), exposing significant cross-architecture optimization gaps. In contrast, Triton sustains 62-101% of cuBLAS performance across all tested platforms without architecture-specific tuning, demonstrating substantially stronger portability.

Abstract

GPU Kernel Optimization

Explore similar work

Evaluating CUDA Tile for AI Workloads on Hopper and Blackwell GPUs

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

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