Source-linked AI summary
The Art of Closed-Formula Defaults: Search-Free Code Generation for Tensor Operators
Paolo D'Alberto, Ashish Sirasao
TL;DR
Materializing the attention score matrix requires O(L^2) memory and traffic that grows quadratically with sequence length. The paper derives GPU tile sizes analytically from operator and hardware descriptions, reporting faster fused attention with less HBM traffic.
Problem
Materializing the attention score matrix QKT requires O(L^2) memory and traffic that grows quadratically with sequence length.
Source: §6 · p. 7Method
Results
1.50× over unfused attention at L=16384, d=128, with 129× less HBM traffic; speedup reaches 1.68× at L=32768.
Source: §8 · p. 11Takeaways & Limitations
The CU-centric derivation yields defaults that update when the hardware descriptor changes, without re-derivation.
Source: §8 · p. 11Fusing the softmax update and second GEMM at register level to eliminate an LDS roundtrip is left as future work.
Source: §6.2 · p. 10
Abstract
from arXiv · showhide
Agentic search and automated optimization of GPU kernels are powerful tools for large language model inference. Their effectiveness, however, depends not on the sophistication of the search itself, but on the clarity of the optimization problem being solved. We provide an application-first approach that drives a hierarchical code generation tool from operator specifi cation down to GPU instructions, and show that a clearly defined computational model makes the optimization problem tractable.
1 Introduction
The paper replaces search-driven tuning with analytically derived, hardware-aware defaults for four tensor operators. It reports GPU measurements validating the derivation and argues that precise operator and hardware descriptions make the method extensible.
1 Introduction
- LDS capacity fixes the CU tile, VGPR budget fixes the warp tile, and CU count bounds minimum parallelism for four tensor operators.The operators are GEMM, 2D convolution, softmax, and fused attention; the tiles are determined without search.Source: §1 · p. 2
- A hierarchical computational model separates operator description from hardware mapping, while a minimal hardware descriptor encodes essential target parameters.The descriptor includes LDS capacity, vector register budget, MFMA instruction shape, and peak bandwidth.Source: §1 · pp. 1–2
- The FeasibilityVisitor checks VGPR and LDS constraints at each tree node and derives tile Q max and tile KV max parametrically.The derivation applies to GEMM, 2D convolution, softmax, and fused attention.Source: §1 · p. 2
- On an AMD CDNA GPU, GEMM reached 12.3 TFLOPS (53% of MFMA peak), while fused attention achieved 1.50× the unfused baseline at L=16384 with 129× less HBM traffic.Source: §1 · p. 2
- The authors argue that precise mathematical descriptions of complex functions and portable hardware mappings could extend the methodology beyond the studied operators and GPU.They describe the derivation as mechanical once those descriptions are available.Source: §1 · p. 2
2 Computational Model
Operators are represented as hierarchical algorithm trees whose levels map computation onto GPU hardware. Analytic tile policies make the tree feasible while keeping boundary handling explicit.
2 Computational Model
- A hierarchical algorithm tree maps device-wide work, compute-unit work, wavefront work, and hardware instructions to successive computation levels.A visitor traverses the tree lazily to execute computation for validation or emit code for deployment.Source: §2 · p. 2
- When dimensions are not tile multiples, each final tile is clipped to the valid matrix extent, and emitted code predicates out-of-bounds accesses.Zero padding is a no-op for accumulation, separating tile sizing from edge handling.Source: §2 · p. 3
- The tree is feasible when tiles fill CU shared memory, stay within per-thread register limits, and meet a minimum active-CU threshold.The corresponding kernel is described as correct and deployable without further search.Source: §2 · p. 3
- For the 4096×4096 GEMM example, each tree level is a tiled partition, and the visitor can stop at a chosen level while the runtime handles remaining levels.The ComputeVisitor validates and the EmitVisitor generates the GPU kernel.Source: §2 · p. 3
40963 GEMM,
The GEMM example uses analytically derived tile sizes, with visitors supporting validation and kernel generation.
40963 GEMM,
- Tile sizes at each level of the GEMM tree are derived analytically from the hardware descriptor.The ComputeVisitor validates at a chosen level, while the EmitVisitor generates the GPU kernel.Source: 40963 GEMM · p. 3
- The figure is identified as an algorithm tree for the GEMM example.Source: Figure 1 · p. 3
3 Hardware Model
The hardware model encodes GPU parameters used in tile derivation. The paper also identifies an MFMA fragment-layout pitfall: instruction shape alone does not specify output ownership.
3 Hardware Model
- The validation target has 64 KB LDS per CU, 256 VGPRs per thread, 64-thread wavefronts, 120 CUs, and 23.1 TFLOPS MFMA compute.The target also has 1.2 TB/s HBM bandwidth and 11.5 TFLOPS scalar single-precision compute.Source: §3 · p. 3
- HardwareDescriptor and MFMAAtom encode the hardware parameters and instruction shapes that appear in tile derivation.The interface is intentionally limited to quantities needed by the derivation.Source: §3 · p. 4
3.1 A Pitfall for the Code Generator
- MFMA instruction shape alone cannot determine how output elements are distributed across threads.The fragment ownership must be known separately by the code generator.Source: §3.1 · p. 4
- On the tested CDNA GPU, the mfma f32 16x16x4f32 output layout transposes the natural ownership expectation, as a structured test revealed.With Q[i, :]=i and K=I, the actual output revealed the transposition.Source: §3.1 · p. 4
- Chained matrix multiplies require either a shared-memory transpose or indexing the next instruction from the confirmed prior fragment layout, not nominal shape alone.The authors advise treating fragment ownership as an empirical hardware property and encoding it explicitly.Source: §3.1 · pp. 4–5
4 Tile Derivation
Tile derivation is CU-centric: LDS and VGPR constraints determine CU and warp tiles, while CU count checks occupancy. The derivation also checks atom alignment and feasibility.
4 Tile Derivation
- The CU is the anchor: satisfy its LDS and VGPR constraints, then check occupancy outward and atom alignment downward.The runtime traverses the tree top-down, but tile derivation builds it CU-first.Source: §4 · p. 5
- For the CDNA GPU, a 64 KB LDS budget with 4-byte elements yields a CU tile of tM = tN = 128 and tK = 32.The CU tile uses LDS for simultaneous A and B tiles, with tM = tN.Source: §4 · p. 5
- The warp tile must fit the VGPR budget: the (32, 32, 8) tile uses 16 accumulator floats per thread, within the 256-VGPR budget.The largest feasible warp tile is aligned to the MFMA atom dimensions.Source: §4 · p. 5
- For a 4096 × 4096 problem with 128 × 128 CU tiles, 1,024 blocks cover 120 CUs, or 8.5 blocks per CU; smaller tiles are used if occupancy falls below one block per CU.The CU count bounds minimum parallelism.Source: §4 · p. 5
- The FeasibilityVisitor checks operator and descriptor constraints, returning tile sizes or the binding constraint; fused attention tightens the VGPR limit because chained GEMMs keep intermediate arrays live.Tile dimensions must also align to the atom shape (16, 16, 4).Source: §4 · pp. 5–6
5 Operators
The paper describes four tensor operators as algorithm trees, derives tile policies analytically, and validates the resulting kernels on a CDNA GPU. The operator examples cover GEMM, direct 2D convolution, and online softmax.
5 Operators
- The four operators are presented in increasing complexity as algorithm trees, with analytic tile policies and kernels validated against a double-precision reference on CDNA.The operators are GEMM, 2D convolution, softmax, and fused attention.Source: §5 · p. 6
5.1 GEMM
- FlyDSL traverses the structured algorithm tree to emit GPU kernels; the same machine-readable specification can feed other backends or automated exploration.The tree serves as both the closed-formula default and the solution space for further exploration.Source: §5.1 · p. 6
- 12.3 TFLOPS for GEMM is 53% of MFMA peak; tiled 3×3 convolution achieves 1.7× over the naive baseline.The GEMM result is validated at machine precision against a double-precision reference, and all 42 convolution algorithm-family tests pass.Source: §5.1 · p. 6, Table 1 · p. 6, §5.2 · p. 7
5.2 2D Convolution
- Direct convolution applies each filter to its input patch without rearranging data; deployment-time filter pre-packing is its stated layout choice.The direct path is selected as the most general approach and exposes the tiling structure clearly.Source: §5.2 · pp. 6–7
- For small input-channel counts such as RGB's IC=3, spatial-to-channel replication to IC=12 reduces partial sums from KH×KW to 1 while maintaining correctness.The described pixel-shuffle trick is used for this case.Source: §5.2 · p. 7
5.3 Online Softmax
- Online softmax keeps a running maximum and denominator in registers, correcting the denominator when a new tile maximum exceeds the previous maximum.Out-of-bounds elements require −∞ padding, approximated as −1038 in single precision; zero padding requires an explicit mask to avoid spurious denominator weight.Source: §5.3 · p. 7
- The online-softmax merge operator is associative, enabling split-N parallelism through binary reduction.The method is validated at machine precision on CDNA at 512×512.Source: §5.3 · p. 7
6 Fused Attention
Fused attention streams key and value tiles while keeping attention scores out of HBM, avoiding the full L × L score matrix. On CDNA, the MFMA+MFMA kernel improves on unfused attention, with speedup increasing at longer sequence lengths.
6 Fused Attention
- Associativity lets each query tile stream K and V while avoiding the L × L score matrix, reducing HBM traffic from O(L2) to O(L · d).The savings scale linearly with L and are independent of tile sizes.Source: §6 · pp. 7–8, Figure 2 · p. 8
6.1 MFMA Kernel
- The MFMA+MFMA kernel uses MFMA for both GEMMs by storing the intermediate Et transposed in LDS; both kernels are validated at machine precision.The MFMA+FMA variant instead uses scalar FMA operations for the second GEMM, and the transposed layout supports reading Et as MFMA A-fragments.Source: §6.1 · pp. 8–9
6.2 Results
- At L≈4096, the crossover separates LDS write overhead that exceeds HBM savings from longer sequences where HBM savings dominate and speedup grows.The unfused baseline uses three tiled GEMM kernels and writes the L×L attention matrix to HBM between steps.Source: §6.2 · p. 9
- 1.50× over unfused at L=16384, with 129× less HBM traffic; fused traffic is 34 MB versus 4,329 MB unfused.Table 3 compares MFMA+MFMA fused attention with the unfused baseline at d=128, tQ=32, and tKV=64 on the CDNA GPU.Source: Table 3 · p. 10, §6.2 · p. 10
- Speedup rises from 1.25× at L=8192 to 1.68× at L=32768, consistent with the L/(4d) traffic-savings formula.At L=32768 and d=128, avoiding the attention matrix prevents writing and reading a 4 GB matrix.Source: §6.2 · p. 10
7 Related Work
The paper situates closed-formula defaults alongside autotuning and analytical tile sizing, presenting them as a correct starting point for search rather than a replacement for kernel optimization. Its contribution is tile derivation for fused operators and documentation of an MFMA layout issue, not a new attention algorithm.
Fused attention
- Fused attention [5] introduced online softmax, while its follow-up [4] added GEMM scheduling improvements; this paper contributes analytical tile derivation and MFMA layout documentation, not a new algorithm.The paper distinguishes its contribution from the fused-attention algorithm itself.Source: §7 · p. 10
Autotuning
- Triton, TVM, and Halide search tile configurations; closed-formula defaults provide a correct starting point that reduces search space at zero search cost.The defaults are presented as complementary to autotuning.Source: §7 · p. 10
Hand-optimized AMD GPU kernels
- HipKittens uses asynchronous loads/stores and multi-wave ping-pong scheduling to address the gap between 53% MFMA utilization and theoretical peak.The paper suggests its tile sizes could start such hand-optimized implementations, separating tile derivation from instruction scheduling.Source: §7 · p. 10
Optimized BLAS libraries
- Unlike CPU SIMD blocking in OpenBLAS, GPU tensor-instruction tiling must handle more complex register layouts, including the documented chained-MFMA C-fragment discrepancy.The paper extends analytical blocking to GPU tensor instructions and fused operators such as attention.Source: §7 · p. 10
8 Conclusion
A CU-centric derivation sets tile sizes for GEMM, convolution, softmax, and fused attention from LDS, VGPR, and CU constraints. CDNA attention measurements show lower HBM traffic and speedup over the unfused baseline, while future work includes broader operators and targets.
8 Conclusion
- The FeasibilityVisitor derives tiles for GEMM, 2D convolution, softmax, and fused attention from LDS, VGPR, and CU constraints, updating them when the hardware descriptor changes.LDS determines the CU tile, VGPRs determine the warp tile, and CUs bound minimum problem granularity.Source: §8 · p. 11
- 1.50× over unfused at L=16384, d=128, with 129× less HBM traffic; speedup reaches 1.68× at L=32768 on the CDNA GPU.The paper also reports a corrected MFMA register-layout formula and a minimal reproducible test.Source: §8 · p. 11
- The authors expect the approach to generalize to hierarchical computation trees on architectures with an essential hardware descriptor.This is stated as an expectation beyond the operators and GPU studied.Source: §8 · p. 11
- Future work includes eliminating the Et LDS roundtrip and extending evaluation to multi-head batching, more operators, and additional GPU targets.Eliminating the roundtrip is aimed at closing the gap toward the roofline.Source: §8 · p. 11