FlashAttention on a Napkin: A Diagrammatic Approach to Deep Learning IO-Awareness
Authors: Vincent Abbott
Gioele Zardini
Venue: Accepted by TMLR ()

Abstract: Optimizing deep learning algorithms currently requires slow, manual derivation, potentially leaving much performance untapped. Methods like FlashAttention have achieved a x6 performance improvement over native PyTorch by avoiding unnecessary data transfers, but required three iterations over three years to be developed. Automated compiled methods have consistently lagged behind. This paper extends Neural Circuit Diagrams for deep learning models to consider resource usage and the distribution of tasks across a GPU hierarchy. We show how diagrams can use simple relabellings to derive high-level streaming and tiling optimization strategies along with performance models. We show how this high-level performance model allows the effects of quantization and multi-level GPU hierarchies to be readily considered. We develop a methodology for representing intermediate-level pseudocode with diagrams, allowing hardware-aware algorithms to be derived step-by-step. Finally, we show how our methodology can be used to better understand existing techniques like FlashAttention. This work uses a theoretical framework to link assumptions about GPU behaviour to claims about performance. We aim to lay the groundwork for a scientific approach to GPU optimization where experiments can address clear hypotheses rather than post-hoc rationalizations.

Giocle Zardini Published in Transactions on Machine Learning Research (03/2025)
FlashAttention on a Napkin: A Diagrammatic Approach to Deep Learning 10-Awareness
Vincent Abbott
Vincent.Abbott.2. 24 @ucc. uk
Deportment of Computer Science, University College London
gzardini Cectiiii@mmmt.oo edu
Laboratory for Information and Decision Systems, Massachusetts Insfitude of Technology
Reviewed on OpenReview: https://peenrenieu..n//// forrmm/mmmmeettottoee pddp AHaA
Abstract
Optimizing deep learning algorithms currently requires slow, manual derivation, potentially leaving much performance untapped. Methods like FlashAutention have achieved a x6 performance improvement over native PyTorch by avoiding unnecessary data transfers, but required three iterations over three years be developed. Autonated compiled methods have consistently lagged behind. This paper extends Neural Circuit Diagrams for deep learning models to consider resource usage and the distribution of lasks across GPU hierarchy. We show how diagrams can use simple relabellings to derive high-level streaming and tiling optimization strategies along with performance models. We show how this high- level performance model allows the effects of quantization and multi-level GPU hierarchies to be readily considered. We develop a methodology for representing intermeddiate-level pseudocode with diagrams, allowing hardware-aware algorithms be derived step-by be--by.-tep. Finally, show how our methodology like FlashAttention. This work uses theoretical framework to link assumptions about GPU behaviour to claims about performance. We aim to lay the groundwork for scientific approach to GPU optimization where experiments can address clear hypotheses rather than post-hoc rationalizations.
be used to better understand existing techniques
1 Introduction
1.1 Background
To exccute operation, graphical processing units (GPUs) must move data from high-level DRAM to low-level compute cores. GPUs are limited much by GB/s of memory bandwidth as TFLOPs of avail- able compute. However, AT models have passed the memory wall algorithms are increasingly limited by bandwidth/ ttaansffr costs (Ootomo & Yokota, 2023 Ivanov et al. 2021 Gholami et al. 2024], as compute capability has improved far more quickly x3/2g than DRAM bandwidth x1.6 2y (Gholami et al. 2024]. Fur-
thermore, DRAM already accounts for 46% of total system power (Ghose et al. 2018). As memory becomes increasingly inefficient relative compute, the importance of considering transfer costs Soats-OO-awarness. et al. 2022 Aggarwal & Vitter 1988 will become even more critical.
DDo
FlashAttention (Dao et al. 2022
Dao 2023 Shah et al. 2024 is an 10-aware approach to attention that
overcomes the memory wall. Attention Vaswani et al. 2017) is central to generative models, including large
language models (LLMs) Mistral AI Leam 2024 Llama team 2024 and image generation algorithms (Ho) et al., 2020 Esser et al. 2024 Rombach et al.] 2022 Podell et al.) 2024). FlashAttention fuses the steps of attention. It computes all soquential steps on low-level memory, avoiding unnocessary intermediate data transfers. It achieves x6 increase in throughput compared to standard PyTorch, arguably making large contemporary models possible.

---

Published in Transactions on Machine Learning Research (03/2025)
However, the conditions under which fusion is possible are not generally exploited. Simple cases like element- wise functions can be compiled into matrix multiplications PaPas et al. 2019 Sabne
2020), but the
bespoke techniques of FlashAttention required manual derivation and three iterations over three years to take full advantage of Hopper hardware (NVIDIA 2022) features. Triton (Tillet et al.
2019) offers some
compilation ion for hardware features but has lagged behind new Alaenttttnttonn algorithms (Da 2023 Shah
ct al. 2024). The current best technique for generating IO-aware algorithms that exploit hardware features remains slow, manual derivation.
Innovating optimized algorithms is essential to efficient model deployment. In addition to FlashAtten- tion, methods like grouped query attention (Ainslic et al.] 2023), KV-caching (Shazeer 2019 and quantiza- tion (Frnntar et al. 2023; Gholami et al. 2022) all reduce transfer costs while having minimal impact on the function we implement or the quality of model outputs. Much like fusion, the success of these approaches relies on understanding the compositional structure of algorithms so that similar but less costly algorithms can be executed. A systematic approach to innovating optimized algorithms will roquire mechanism for understanding the compositional structure of algorithms along with performance model which compares varying means of executing the same operation.
The hardware characteristics of GPUs have a significant impact on performance which varies depending on the target algorithm. When choosing between A A100s, H100 SXM5s, or H100 PCTes (NVIDIA) 2022 p.39), we must consider the varying compute, bandwidth, intermediate hierarchies, architecture features, and low- level memory, for which we pay in environmental and economic resources. The application of these features is often non-obvious, FlashAttention-2 (Dao 2023) was released while the hardware for F already existed (Shah et al., 2024], which achieved 75% improvement in forward speed. Understanding the impact of GPU features is necessary component of innovating optimized approaches and making full use of deployed resources.
1.2 Contributions
This paper contributes representational scheme for deep learning algorithms based on Neural Circuit Diagrams that shows the distribution of tasks across GPU hierarchy and associated resource usages (Section 2.. This scheme incorporates theorems about the compositional properties of fused algorithms (Appendix A.I., allowing for the rapid derivation of high-level sketches of GPU optimized matrix multiplication and attention, along with performance models (Section 3). This performance model can consider quantization and multi-level hicrarchies (Section 4 bolstcred by thoorems (Appendix A.2). With a model corresponding to Hopper-like ure,itecture, we use step-by-step process to derive hardware-aware algorithms (Section 5 These derivations consider coalesced memory access, tensor-core operations, and overlapping operations, and reveal the degrees of freedom in the algorithm's configuration. Our methodology can be used to analyse the bottlenecks of existing methods like FlashAttention (Section 6/.
The main value of this work is providing a potential framework for relating assumptions about GPU behavior to claims about performance. This allows future empirical work to address and improve specific assumptions, thereby iteratively improving the model. This contrasts with the prevailing "tinkering" approach within deep learning that develops post-hoe intuition based on past experiments. Post-hoc intuition is trained the "test set" of prior successful approaches, weakening gencralizability and scientific value. Our diagrammatic scheme allows the rapid application of the framework to algorithms used in practice distinguishing it from descriptive theory and toy model approaches. Diagrams leverage human visual comprehension and enable complex thcorems and underlying implications to be encapsulated with simple notation.
2

---

Published in Transactions on Machine Learning Research (03/2025)
2 Diagramming Deep Learning Algorithms
2.1 Diagramming Functions and Data Types
Diagrams have alternating columns of data types and functions. Data type columns are shown in Figure Arrays such as Raxbxe are represented by a wire for each axis labeled with the size. Data types may be tuples of arrays, such as Raaxbx x idxxexe and are represented by placing a dashed line between constituent arrays.
Figure 1: We represent arrays, of forms such as Raxbxe, by labeling stacked wires in column To represent data types that
The diagram represents the array type Raxbxe by stacking the labeled wires. b
with b, and consist of lists of arrays, such as R= gaxbxe x gdxe we place a dashed line between them.
C
d Stacking with a dashed line assembles the data type Raxboe x R
Functions between data lypes are represented by labeled boxes or pictograms with their input/outpul shapes the left/righ /right. Sequential execution (composition) of functions is shown by horizontal placement (Figure crcating diagram with alternating columns representing data types and functions. Parallel exccution (concalenation) of functions stacks them with dashed line in between (Figure 3). A concalenated function takes concatenated inputs and provides concatenated outputs. The change in the input/output is reflected by the diagram.
Functions are represented by la- beled boxes or pictograms, which aid intuition. These representations posed, which represents sequential execution and yields diagram with alternating data type and function columns. We represent. composition by F; F.F.
Figure 2:
The function F has the form
be horizontally
R
x bxc - RR x R
Horizontal placement
a
a
b
composes functions,
b
F
G
that this block gives
e:
e
C
R: - bxe
Functions can also be stacked with separating dashed line, which concatenates their inputs and outputs. Concatenating tuples @ is considered to be associative. For concatenated functions F 1 G, if (()) x' and G(y) = y/ then (F (I & F(x) G(y).
Figure 3:
We can concatenate functions, so
F
we get FOG:R x
lxc x
R x R x automatically shows the new input and output data types.
XXc
The diagram
G
We represent identity functions that leave inputs unchanged by extending the data type. This reflects that composition with identitics leaves function unchanged. Functions are stateless and are defined by how they map inputs to outputs. Therefore, we concatenate with identities to represent function acting on some data but leaving the rest unchanged and available. With these tools, we can build compound diagrams such as Figure
3

---

Published in Transactions on Machine Learning Research (03/2025)
A compound diagram can be dis- assembled into columns representing alternating functions and data lypes. Stacked functions and data types can be further decomposed to find the core units concatenated to construct them. Iden- tities are represented by continuing the represen- tation of data types.
Figure 4:
We have alternating columns ofidata types and functions These columns concatenated along dashed lines.
Continued wires in function columns are identitios.
F
H
Functions be mapped over additional axis, represented by weaving the axis the function's outputs and choice of the inputs (Figure 5). This diagrammatic implementation naturally updates the input and output sizes for the mapped function. When input segment is not weaved, its data is copied to evaluate each index along the outputs of the new axis. The axis can be weaved into any location of the target segments.
SoftMax on a vector R
SoftMax on rous of R
Contraction for vectors R
Matrix Multiplication of Raxb IR bxc
a a
X
b
x
X
b
b
b
C
C
Figure 5: A function can be weaved, which adds an axis to the outputs and some of the inputs. The function is mapped over this axis. When we weave the "item" I array represented by thick dotted wire, we can remove it Here, we provide a weaving for SoftMax, represented by triangle, to have it act over each row of an array, and of linear contraction (dot/inner product), which provides matrix multiplication.
Weaving function allows for complex mappings to be represented and avoids the ambiguity of typical expressions. We can weave primitives defined itemss such multiplication, addition, and copying, We use weaving to represent the splitting and joining of shared axes, which overcomes the typical difficulties of expressing how arrays are partitioned and concatenated. We show this in Figure
Addition IR x R R
Multiplication R RxR-E
Copying
Splitting vector
Joining Concatenate vectors.
R-RXR
R"
R
R k
R" x R 3+
n
+ m
Addition of rous
Multiplication
Copying
12
in
with a vector,
Outer product,
Rax array
Splitting
Joining
Zij = Xij + Vj
Zijk = x1 Y
An array olong an aris.
Concatenate arrays.
b
b
b
b
a
b +
b
+ m
a
11 k
m
Figure 6: Wcaving primitive functions let express the addition of vectors to each row of array, multiplication over specific axis, and the copying of arrays. We can use weaving to split array along slices of an axis or show the axes over which arrays are joined.

---

Published in Transactions on Machine Learning Research (03/2025)
2.2 Representing Deep Learning Algorithms
We have so far expressed functions maps between inputs and outputs diagrammatically: Deep learning models employ algorithms, which implement function but have resource usages and inputs/outputs located at different levels. We embed algorithms in hierarchy. A hierarchy consists of levels connected with pipes (as in Figure 7) which allow for memory sharing with family of cores located at the level below. The available algorithms are restricted to those provided at each level of the hierarchy.
H100 SXM5 GPUs consists of 132 streaming multiprocessors (SM), each with compute capabilities. The GPU has global GMEM memory, two 1.2 caches, and shared memory memory/register files SSMEM/RMER per SM. Logically, this is organized (which pipeline instructions) each with 32 threads. Threadblocks have independent allocations of SMEM, and threads have independent allocations of RMEM. We abstract these details away for our hierarchy, focusing on the logical organization and memory sharing capabilities.
grid (device-wide) of threadblocks (limited to an SM) consisting of warps
We can diagram hierarchies using graph showing the available levels and their connections. These hierarchies model real GPU, and pro- vide levels corresponding to logical ab- stractions.
Figure 7:
Abstract Two-Level Model
Abstract A100
H100
Three-Level Model
Tensor- Core
Tensor-Core
Device
GMEM
co
Device
GMEM
Thread block
Distribnted
cluster
SMEM
Threud block
C1
E1
SMEM
Thread
block
SMEM
Tensor
'Thread
{2
Tensor
Thread
Core
Registera
Coro
Registerss
We use colors to represent levels, and color arrays to diagram where they are located. In this section, we use a two-level model where higher level 20 arrays are colored black and lower level arrays are colored orange. This lets us diagram algorithms as in Figure
Figure& 8: Here, we diagram algorithm which takes SoftMax over E array and contracts it over a E x d array. We per- form transfers to move data to lower levels for computation. This diagram shows sequential execution, concatenation, and weaving of algorithms diagrammatically.
Load to Ll Save to 0
9
5
T
We are interested in algorithms' total transfer cost He and maximum memory usage per core Mf (memory usage). These resource usages are defined per level and measured in number of values, and can be determined from diagrams as in Figure Total transfer costs from level, equal to the sum of the size of arrays changing colors. Memory usage is lower bounded by the maximum size of data at level for any column. We aim to minimize the total transfers while keeping memory usage below a hardware limit per level, Mmaax
equal to the total size of data loaded to and saved
The SoftMaxx-Contraciinn algorithm from Figure B can have its transfer cost derived from the total size of data changing colors and its memory usage lower bound determined by the maximum data present the lower level at any point. We assume that E > d.
Figure 9:
M H
7x+xd Xd
Cumulative 4g + Rd
QX
7x
9
+ Rd
qx
q
x
X
5

---

Published in Transactions on Machine Learning Research (03/2025)
As diagrams show all data types, operations, and their arrangement, can adapt our performance model to consider all aspects of an algorithm's performance. Using diagrams, we can approximate compute by taking the compute required to execute an algorithm multiplied by the size of axes it is weaved over. A k-size contraction requires 2k FLOPs; therefore, m x k by k x n matrix multiplication requires 2mkn FLOPs. In Section 5.7
this used find the clock cycles required per column to overlap computation.
2.3 Group Partitioning
The first optimization strategy we introduce is group partitioning (Liling) a mapped algorithm (Figure 10 If an algorithm is weaved over an axis, then the axis can be split, the mapped function applied, and the axis rejoined to yield the same result. Each sub-algorithm acts on a batch of the data in a separate core with reduced low-level memory usage.
A weaved algorithm is functionally equivalent to sub-algorithms acting on partitions of the weaved axis. We use = to indicate fune- tional equivalence, meaning algorithms map the same inputs to outputs but may have distinct resource consumption profiles, and therefore are not strictly equivalent. These partitions can be of any size, which write Ja. We can recur- sively expand the expression on the right until d' S Ja- The unweaved segment of the data must be loaded by each sub-algorithm.
Figure 10:
Theissuballorrthhe are executed
across different processors at the
saine level.
8
a a
a
F
V
y
W
V
207
y
W
We can diagram this strategy by labeling the weaved axis with the target group size 9a while the lower level in Figure 11 The low-level memory usage Mg for the diagram is then calculated using this group size, not the full size of the axis. Each sub-algorithm needs to load and save its batch input and output data. The per-group transfer cost Hr.g is calculated using the In group size for the partitioned axis which is multiplied by Nt.g = a/ a/g, ga batches to attain He- Ne.
Non- Kon-grouped inputs are sent to all active cores at the lower level, meaning their transfer costs are multiplied by Ne., without reduced per-group transfer costs. Smaller group sizes ga decrease memory usage but increase gr increasing He there an unweaved input. To reduce total transfer costs, we must find the maximum 9a value that does not exceed maximum memory usage M.nax.
Figure 11: Group partitioning can be represented by relabeling an algorithm with the partition batch size (group size) at the lower level. The group size is used for memory usage and per-group transfer cost calculations for the lower level. Per-group transfer costs are then multiplied by the number of groups Ne.g = a/ga to give the overall transfer costs.
Mt 8,x w
&aM
Cumulativee
Heg 8,x + w
8ay
8a(x + y) +
Mn > 8o + W, gay
&a
a
Hn = 8x(x + +)) +
x
x
Nn = a/
F
W0
20
Ha = ax + ay atv /
The total number of sub- algorithms is the product of the number of batches for each axis, Nr., ab/ gagb. The relabeled sub-algorithm represents the memory usage and per-group transfer costs of each sub-algorithm, using the group sizes Ja and go for resource usage calculations. We then multiply the transfer costs by Ne.g. We can use this to determine the optimal group sizes for an algorithm grouped over multiple axes, such as matrix multiplication (sce Section 3.1) and to determine whether it worth grouping a small axis or transferring its full size.
If multiple weaves are relabeled, the data is batched over each in Figure 12

---

Published in Transactions on Machine Learning Research (03/2025)
Figure 12: A function with multiple weaves have relabeling applied to cach of its weaves. The relabeled sub-algorithm provides the mem- ory usage and per-group transfer costs using the group sizes.
MÅ1 8 wgg LaMSh
Cumuuatine
H How Sax + ugg
8,X + -8.0 + 8o8by
Mn 7 gax +.20, gaghy
a
a
HAL = NN1 gbt ga8by)
x
X
V
Nug ab ab/8888 gagt
F
2
2
:. H1 = abx +8a + aby
gb
2.4 Stream Partition
Stream partitions (recompulation) exploit recursively decomposable polymorphie functions to feed data in batches while maintaining intermediate outputs on-chip, reducing low-level memory usage. Functions can be streamed if they are polymorphic for a specific axis (defined for that axis being of any size) and have an accumulator that can incorporate incoming data recompute the maintained output, as shown in Figure This allows for recursive expansion (see Figure 14 that maintains minimum data on-chip at any point.
13
The condition for streaming quires that function F be polymorphic along the axis a and can be decomposed in the man- ner above, requiring the existence of another polymorphic function B called the accumula- tor.
Figure 13:
A
Y
Y
X
V
B
F
F
X
X
If the condition in Fig- ure met, the function can be re- cursively decomposed until a' S Sa- This allows the function to be evalu- ated from batches of the input data.
Figure 14:
a
S₁
Y
y
X
1/
B
X
F
F
B
B
If an axis originates from a transfer and is fed to a recursively decomposable polymorphic function, then it can be relabeled with the streaming batch size (stream size) as in Figure 15 This creates a representation of the aub-algorithm B which is repeatedly applied to process the data. We need to add the output size in parentheses the input consider its contribution to memory usage. The memory usage of the algorithm is then determined using the stream size s instead of the full axis size b. As we eventually stream the entire axis, use the full axis size b to evaluate transfer costs. Typically, strictly benefit from limiting the stream size to 1 as this reduces memory usage while imposing no increase in transfer costs.
We can relabel a stremmable axis with the batch size 8, as it is transferred. We are required to add the y output array shape at the input to the streamed algorithm. This lets the relabeled diagram derive the memory usage at the lower level using the stream size 8b- As all data along the axis is transferred to the chip, the full axis size b must be used for transfer costs.
Figure 15:
M as + 1 + V
V/
Cumulative ab + w +
H ab
+ w
y
Mn > aSs + +0 + He ab w +y
b
Sp
F
wo:
w
Per the fusion theorems of Appendix A.1 streamable are resistant. to modifications. The streamable axis may be single axis of an array, and composing or weaving streamable algorithm while maintaining this axis yields a streamable algorithm. This allows the stream labeling to be maintained for resource usage evaluation as in Figure I6 and allows the streamability of complex functions like attention to be derived from a streamable kernel as in Figure 21] In Figure 17 we apply group partitioning to a mapped streamable algorithm. We use 9n for per-group transfer evaluations, and both 9n and 8, to evaluate memory usage.

---

Published in Transactions on Machine Learning Research (03/2025)
For modified streamed algorithm, we can contimue to use the stream batch size for memory usage evaluations. As the function generates x distinct y values, it needs to main- tain each on memory, resulting in y x q.X X r main- tained memory before and after the repeated E algorithm.
Figure 16:
Ma gas, + XT + vgr gas + X7 + yar
Cuncuaatio
Hrg ' b + x'
a'b + x' + 1/4
9
9
a'
a'
E
4
b
55
Sp
5B
F
Y
K
1/
xx
(gyr) : (gyr) j
Figure 17: We can apply multiple relabelings to an algorithm. This lets us find the per-group memory usage and transfer cost. As the func- tion is mapped within each group, it needs to maintain 9q copics of the maintained y data, in- creasing its memory usage.
MA g, + w + 8oM
Cumulative
Hes 8 ab w
g,ab w 88y
9 89
Ma > &gaSs + w + 8qy
q.
Hig .gab + w + gay
b
Sb
F
y
Ne1g =9/8g
w
:.
He = qub + qy + qwo 8a
3 Examples
3.1 Matrix Multiplication
As contraction (dot product) is streamable (see Appendix A.3.1 streamability of matrix multiplication, its weaved form. This provides diagram that supplies a performance model. We then optimize for the batch sizes to minimize total transfers given some maximum lower-level memory usage M.
use it as kernel for deriving the
Matrix Multiplication
Contraction
ga a
N₂ Ja9e
ac
H Ny
b S
b
Sb
= ac
Hg = gab + bge + 9a9c
(gab + bge + 9a9c)
86 b
b
Sb
Ja9c
(1)
&c (ga8c)
M > Ja9c 9asb + sb9e
= abe(ge -1 + 9a
-1 + ac
VM > J. Ja
> 2abe -05 + ac
Figure 18: The dot product is strcamable function. Therefore, matrix multiplication, which is the weaved form of it, is also streamable and can be group partitioned.
This simple derivation reveals an immense deal about matrix multiplication and its associated performance model. Firstly, we have derived that Ja should equal 9c. corresponding a square tiling shown in Figure [8] We see that bandwidth increases with b at rale dependent. on the size of memory. A standard approach would effectively assume that M co. This would cap 9u at and ge at C, not the memory limit, yiclding H ab + be + ac (Gholami et al. 2024 Ootomo & Yokota 2023). Considering the tiling size and memory constraints, we get H 2 2abcM
ac In Appendix B.1 we examine the arithmetic intensity.
Furthermore, we can use matrix multiplication as an instructive preview of the next section which focuses on multi-level performance models. We can compute each ga X 8, x Se matrix multiplication at a lower, green level. As square tiling of ga = 9e optimal for sizes a and c, square tiling of ha he optimal for sizes 9o and 9e- This generates recursive tiling pattern which allows the two-level model generated by simple diagrams like Figure 18 to be extended to complex GPU hicrarchies.
8

---

Published in Transactions on Machine Learning Research (03/2025)
Figure 19: Figure 18 represents a tiling process where we split the a x b matrix into tiles of size Ja X and the b x matrix into tiles of size 86 x Ie- Each lower-level processor is assigned responsibility for a Ja x 9c of the output, for which it accumulates outputs by iterating down Sp. The effect of go and ge on total transfers reflected in the corresponding equations of Figure implying square tiling is optimal.
We iterate down the b axis tiles of size Sb--
gc
a
8
Figure 20: If there is green level below the orange, then the Ia X 9 responsibility tiles are further split into blocks of size ha X hey each assigned to lower-level processor. Each stream loaded to the orange level can be split into up streams for the lower level. Re- sponsibility assignment from the black to orange levels is similar to respousibility assignment from the orange to green levels, allowing us to recursively extend the two-level model. This algorithm is dia- grammed in Figure 24
C
44
b
b
he
u
a
7
3.2 Attention
We derive the streamability of attention from the fusion theorems. We begin with the fact that SoftMax- Contraction is streamable (Appendix A.3.2). Then, we can compose with a contraction over the queries as an E algorithm from Figure [6 This generates a streamable algorithm, which we weave with the 7 and
d axes. This generates Figure 21 Correctness is ensured as the diagram gives the typical expression for attention, O = Max (Q
K k V, with axes clearly indicated. We can then label q to distribute the queries across processors, yielding the FlashAttention algorithm. Figure 21 then, can be seen as deriving and providing a performance model for FlashAttention.
Attention
SoftMaxx-Conzratiion
g
87
Ng = 4/9a
H N Hg
d
d
x ST F
Hg 297d 2rd
(297d + 27d)
X 5y
9
ST
d
d
M > 297d 2syd
27d 27 97
(1)
X Sy
95 & M/2d
> 2nd + 4rgdd M-1 11
d
d
(g
(g
Figure 21: SoftMax followed by a contraction is streamable. We precompose with contraction E weaved by 8, and provide weavings by 7 at the top and d at the bottom to construct attention.
We can apply a similar technique to find the transfer cost of grouped query attention (Ainslie et al. 2023) (Figure 23) and multi-head attention (Vaswani et al.) (017)) (Figure 22 These use additional weaves, but their evaluation remains straightforward. This shows how diagrams can be used to both derive optimizations and experiment with modifications to the algorithm, motivating further innovation.
9

---

Published in Transactions on Machine Learning Research (03/2025)
Figure 22:
Multi-head
Multi-llead Attention
attention conducts multi- ple attention algorithms in parallel. It is represented by an additional axis weaved over every opera- tion. Notice how scaling h leads to linear change in costs, while d leads to quadratic change.
0 8
Ng = gh/(9ggh)
H = NgH
h 8a
qh
d
Hg = 2gggnd + 2nd
(29gthd + 2%d)
949h
Sx
M > 2979nd + 28yd
-2qhd + 2rghd/ (179g)
h gh
d d
Sx
9q9h
& M/22d
2hgd + 4hrrqd M-
x 87
h 8h
h
d d
d
(898,4)
(g98)))
Figure 23: Grouped query attention has additional weave accompanying the queries. This reduces the required parameters and compute to gener- ate the keys and values. However, we observe that the required trans- fers similar to full multi-head attention.
Grouped Query Attention
87
Ng /
H NgHH,
d
Hg = 207 d + 2nd
(2g70ed d + 2nd)
979y
Y Sx
M > 2g + 23
29 + 2g (9g9))
d
87
q
: 949a < M M/2d
2ggd + 4gzqd M-
I ST
d
(g8.4) (898.))
4 Analysis of Performance Models
Once a two-level model optimization of an algorithm is found, we can extend it to consider a multi-level hierarchy. Each lower level has tiles which fit inlo the level above (Figure 24), meaning the optimal strategy and performance model extend in generalizable manner. We can create universal performance model for transfer costs which considers the impact of the GPU hicrarchy and the transfer rate and memory caches at different levels. This allows to make informed choices between different GPU architectures given their energy and capital costs, levels of quantization employ, and the configuration of GPU hierarchies.
4.1 Optimal Transfers H (d,
Applying the two-layer model to diagrams provides optimal transfer costs H1" H1" M)) given some configuration of axis sizes ū and lower-level memory M. So far, these expressions have a standard form given by the sum of power functions which solves for Equation in Appendix A.2
H² Σ (7) M
(1)
The index iterales over lerms, the coefficient (x) is dependent on the axis sizes, and 3 > 0 is an exponent greater than zero as transfers necessarily decrease with increased memory size. The exponents 3, indicate the scnsitivity of performance to memory size and indicate how data distributed. For attention, the -1 factor indicates the data broadcast all groups, while for matrix multiplication the M A factor indicates square tile distribution.
10

---

Published in Transactions on Machine Learning Research (03/2025)
4.2 Multi-Level Performance Models
An algorithm with multiple levels requires H* (", Mr) transfers for each. Even though data cannot be directly transferred from the highest to lowest levels, lower levels can utilize the data loaded and saved by intermediate levels. The execution of H* (ă, Ma) intermodiate level transfers makes data available for the lower level E and accounts for saving back. Each intermediate-level tile fits a larger number of low-level tiles (see Figure 20), among which it distributes weighted transfer cost Hi," each level. For the highest level, we assume Mo + oo and 1 = 0, data is already present. This means that the total weighted transfer cost of algorithm can be expressed by:
data. This fitting process has a negligible error with large M. We assign
H' = - Ha( = Ea(a) a(ā) Ehe' Ma
(2)
Therefore, the relative performance of GPUs is determined not just by the raw transfer rates but also by the memory size of available levels and the specific algorithm being implemented. For attention, the key factor per level is ~ H-1/M while matrix multiplication has H-1/M This implies that atlention is more sensitive to the size of lower-level memory and that less lower-level memory (smaller tile sizes) is required to avoid bandwidth constraints.
4.3 Quantization
Equation 2 and the two-level model consider transfers and storage limits in terms of number of values. However, GPUs are restricted by the number of byles we can transfer and store. If we have q byles per value, then the maximum number of values Mr. M MBBtts /q and the transfer weight is H. (H, (A7Byess) /q) Substituting these expressions into the total transfer cost, we get:
Hy = E /q) H ( MBytes/s)
Ea(a) Eiir (MByttogy- -B g1
(3)
As 1+33 >> 1, total transfers are superlinear to the degree of quantization. Halving the quantization from FP32 to FP16 can accelerate attention by up to x4, and improves large matrix multiplication by a factor of 21.5 s 2.83. This indicates that a generous use of quantization and specialized kernels critical to high-throughput implementations of models.
4.4 Intermediate Caching
We can choose to store output data at lower levels, and save it up in chunks. This changes the level immediately above the lower level to caching level, which we indicate by adding asterisks to its output data as in Figure? 24 The size of this column not memory restricted by the intermediate level which is only used to temporarily store data as is sent up. However, the lower levels must remain active to store data, and this imposes the restriction that N Ng.22/Ng.11 & Nmax which is hardware limit. With oulput restricted algorithm, this results in H² (ā, N M22) transfers being required for the intermediate level fl, using the total lower level memory N maxx Mr2 instead of its own hardware maximum memory M This is elaborated in Appendix A.2.1
4.5 Cross Transfer Levels
Our model can encompass levels that perform inter-core communication instead of providing shared memory by using modified weighted transfer weights. These cross-transfer levels encompass H100 thread block clusters (Luo et al., 2024), multi-GPU systems, and intra-warp communication. We set up as a higher level, as a cross-transfer level, and e as a child level composed of linked processors. Instead of sending H* (ā, M.) data directly to children, we send H" (7, M.Nmax) data to any of the interconnected children and cross-trance transfr
11

---

Published in Transactions on Machine Learning Research (03/2025)
With multiple levels in hierarchy, we can maintain output or streamed data at the lowest levels and use the intermediate levels as a cache. The cache size does not contribute to the 21 max restriction. Instead, we require that Ne2.g/Ne, & Nmax If an algorithm is output restricted, can implement this by the cache size being less than Mj- M
Figure 24:
MA 8aSb + 5b8c
Ma h1s + hug. + h,h,
hahe
Cummulatioe
Hos g.b bgci
8a80i
8.b + bg. + 8o&e
He2g
Hn, bitel
(halt
h1b + blt hahe
a
8a
hg
A
8o80 888.
M
max > ht
Sh
Sh
Nmax M
max 8age
b
Sb.
Sh
Nag /Ni & max
h.
(h,
the remaining data as in Figure 25 and levels, adding a level x between h and c with transfer weight H - H and memory M.N and replacing the transfer weight of level e with Hzte- We outline this derivation the Appendix A.2.2
This results in a performance model with modified transfer weights
To send data to children, we directly transfer (ū, MoNtmax) distributed across the child processors, treating the cross- transfer level as having memory of size M We then perform the remaining H" (a, M.) H Hā MoNe transfers as fast cross- coos-transsers.
Figure 25:
Cross- Transfer GPU
Perspective of child process
h
a
h
5 Pseudocode and Hardware Optimizations
So far, we have focused on abstract models that represent the generic assignment of tasks across GPU hierarchy and investigated the accompanying performance model. This systematic methodology provides the same level of detail as the theorems from FlashAttention (Dao et al 2022).
The aim of the diagrammatic approach is to quickly develop optimized kernels. The success of F (Dao 2023 and FlestAttem (Shah el. al., (024) relied on considering specialized hardware features; tensor core operations and asynchronous compute (in the case of Hopper). Standard methods required years of exploration to incorporate these aspects. With diagrams, we can systematically identify the contexts in which hardware features can be applied and quickly derive kernels.
In this section, we develop the tools to go from generic, abstract analysis to developing an algorithm for specific hardware architecture. In this section, we provide systematic procedure to go from abstract two-level model to an algorithm configured for a specific GPU architecture. We will work with a toy hierarchy which imitates Hopper, with levels for the different modes in which memory can be stored; SMEM, registers, or tensor corcs. Available subalgorithms will be dictated by the functionalitics of the respective levels. The aim of this section to not focus on Hopper specifically, as many aspects of implementation will be missed. Instead, it is to show how diagrams can be used to systematically derive a hardware-aware algorithm. This methodology can be extended to Ampere, Blackwell, and non-NVIDIA architectures.
Coalesced Memory Transfer
Between the device and SMEM memory, GPUs move data in chunks of 128B of consecutive memory. This is remarkably straightforward to represent with diagrams. Arrays represent how data is distributed across each stride, so the lowermost axis of an array represents consocutively stored memory. If we enforce that the lowermost axis is divisible by 128B/q when assembled in the device and transferred between GMEM and SMEM, then we can assure coalesced memory access. This may require that each thread block stream loads
12

---

Published in Transactions on Machine Learning Research (03/2025)
data for multiple lower-level streams. If using SMEM as cache, there is usually plentiful memory available for larger streams.
A floating divisor in the superscript of an axis/batch size is used to indicate a value it is divisible by (see Figure 26 This is done at the point where the restriction imposed and along the immediately weaved axis. Multiple divisors impose the least. common multiple.
Tensor Core Operations
Tensor cores provide very last matrix multiplication, and their consideration was the major contribution of Attennontt (Dao, 2023). Modern GPUs have far more FL.OPs available for tensor cores than general-purpose register operations (NVIDIA) 2020 2022). However, tensor cores require memory to be fragmented in an incoherent manner across warps groups of 32 thrcads) warpgroups (groups of 128 threads, Hopper only) which makes data unsuitable for general-purpose operations. We can still perform element-wise operations, as these do not require information about the location of data. As result, tensor cores act like distinct level of the memory hierarchy, having access to a different set of algorithms than general-purose thread data. Coherent SMEM data can fragmented for tensor cores, providing a pipe between the threadblock and tensor core levels.
Tensor cores (wmma or #gmma) can only manage data at certain sizes and quantizations (NVIDIA 2024) (mmma, Wgmma). Quantization be implemented by adding multiple smaller matrix multiplications, making divisibility the critical factor. We can enforce this restriction by placing superscripts for tensor core axes, as shown in Figure 26
be represented by floating tag. Matrix multiplications of larger sizes
Figure 26: Multi-level matrix multiplication uses the SMEM level to cache data for lower-level tensor core operations. We enforce the divisibil- ily restrictions for coalesced SMEM transfers and tensor cores using superscripts.
PP16
FP16 8
8
(64) (641
11
(16)
PP16
b
Sp
(16)
c (64)
g(
(64)
8.
c(64)
(hght)
5.1 From Diagrams to Pseudocode
We can expand streamed algorithms into looped pseudocode forms where all variables are explicitly shown as in Figure 27 The columns pseudocode diagrams provide the size of variables required in memory and the transfers/operation we need to apply. This allows us to pre-allocate memory to derive the exact memory usages, as well as per-group transfer and compute costs. Columns act like lines of code but more clearly express the details of axes and available optimization strategies than polymorphic streamed algorithms are defined for the stream axis being of any size, can begin the algorithm with head F taking an axis of size 0, initializing the maintained output to incorporate further information. As this expansion is built from the recursively expanded definition of a streamable function, correctness is ensured.
methods. As
Stream
S
a'
b
E
a' a'
E
b
Sq
b
Sh
F
K
x
G
B
B
y
K
x'
x
G
Deletion
Figure 27: A streamed algorithm can be re-expressed with an explicit loop. The diagram illustrates how axes are partitioned, which variables are maintained, and what operations are applied iteratively.
13

---

Published in Transactions on Machine Learning Research (03/2025)
By substituting the subalgorithms in Figure 27 with operations relevant to particular algorithm, develop detailed pseudocode expansions for any streamed algorithm. This expansion is the first step in developing hardware-aware algorithm. SoftMax--ontractio,, for example, can be expanded by substituting B for the contraincont accumulator from Appendix A.3.2 shown in Figure 28
Stream
8
B
8x
1
/2
Figure 28: An example of pscudocode expansion. We re-express SoftMax-contrration using the accumulator derived in Appendix A.3.2
Pseudocode expansion applied to attention involves substituting with query-key matrix multiplication, F with SoftMaxx-cntraation and, therefore, with the accumulator of SoftMax-contraction. We are also required weave the top (which is partitioned) and the bottom axes, deriving the two-level model pseudocode expansion for attention in systematic manner. The outcome of this process is shown in Figure which provides two-level model expansion.
29
Stream F S3
8
Deletion
8
R7
B
d
S
84
X
87
8
8q
d
87
388
Init.
8q
8a
d
Step I of deriving a hardware-aware algorithm is expanding the two-level model into pseudocode expression. We can apply the fusion theorems to attach E and weave over the ş and d axes.
Figure 29:
5.2 Subloops
Pseudocode expansions allow for streamable and split subalgorithms to be identified, which allow incrcased flexibility. Streamable subalgorithms are identified by inner parentheses (subloops). Note how accumulators are streamable, and therefore, the B component of Figure 27 always allows for a subloop that streams s, along chunks of size uq. Furthermore, matrix multiply-add operations can be split along any axis and linearly accumulated. We represent this by placing dotted box around the operation. Shown in Figure 30 this is the second step in systematically developing hardware-aware algorithm. At this point, still have hardware-inddeenndent two-level abstract representation of the algorithm.
14

---

Published in Transactions on Machine Learning Research (03/2025)
Stream X 7
7
84
Deletion
d
d
Split
X
8q
84
81
R4
Spli d'
FP16
8
d
d
3 89
Init.
87
1z d
Stream Sg - Mg
Figure 30: Step 2 requires identifying the subalgorithmms and, hence, the degrees of freedom in axis sizes. 5.3 Recoloring Diagrams to Utilize Tensor Cores
We can impose hardware features by recoloring matrix multiplication and associated data as blue, and general-purpose operations as green. Tensor cores and thread levels are partitioned into groups within the 97 SMEM blocks, and we relabel 9T to la for threads and wy for warpgroups. Transferring between these levels requires piping through orange SMEM. We also impose quantization labels. This generates Figure 31 At this point, the diagram only assumes the presence of tensor cores, and is therefore applicable to Volta, Ampere, Hopper, and similar non-NVIDIA architectures. Note that the second matrix-multiply add involves scaling by 8 weaved over the Dt axis. This is not possible with tensor cores, without complex diagonal matrix approach which falls outside the standard paradigm of a systematic, diagrammatie derivation. This distinguishes our approach from FlashAttention.
Stream X
57 FPS
FP8 7 8q
Wg
Deletion
d
PP8
RP16
Split d
d'
X
8
&
1g
8n
FP16
d
8y
11
M
Spli
x FP16
8x
PP16
14
8
d
d
Init.
RP10
FP16
5
1z d
Stream 83
Figure 31: Step involves assigning operations/data specific cores and quantization levels.
5.4 Applying Divisor Constraints
Next, we impose the divisor constraints specialized hardware features, including coalesced memory access and tensor cores. The divisor constraints come from the specific architecture being used. By imposing that (128) is divisible by 128, we indicate that we are working with warpgroups. We draw the axis sizes for Wy
tensor core operations from the specifications of Hopper
15

---

Published in Transactions on Machine Learning Research (03/2025)
Stream X
- 8x
FP829
7
81
001 (129)
64
(2)
Deletion
PP8
1/166
X
32
64
&
8
Splat (30)
d
32
(
Split
FP16
FP16
d
((64)
((1))
:66
16:
d'dlor
Init.
EP16
Fg
FP16 ga
7
((16)
d' (16)
(1
2 d
Stream 83 - 1z
Figure 32: Step 4 has us impose divisor constraints, configuring the algorithm for a specific architecture. In this case, Hopper.
5.5 Identifying Variables
The data type-columns of diagrams indicate the location and size of variables simultaneously required in memory to execute the algorithm. Performance is best achieved by pre-allocating all required variables. From Figure 33 we identify all the required variable sizes. In some cases, we can reuse variable at two points whe where it is not simultaneously required.
Stream I
FPs 8q
Wq
64
Deletion
d (128)
FP8
RP16
P
Split d(6)
K
32
64
((4)
[64]
184
FP16.
Split d' d"
d (64)
(128)
32
8A
D
AOÃO
y PP16
(22)
MP16
V
64
S lT
(664)
(64)
(36)
16
16:
PP16
Init.
RP16
FP16 ga
d(6))
(16))
((6)
d
Stream 88
Figure 33: Step identifies the required variables for accurate memory assessment and pre-allocation.
5.6 Configuration Table
Our aim is to find a configuration of axis sizes that maximizes performance while conforming to memory constraints. The memory limits are given by Hopper's specilications) which limits the SMEM and register memory per thread block. In Table [ list all the variables, the memory they use, and how they scale with the number of warpgroups, NwG = 97/ This allows us to find the maximum warpgroups per thread block given the configuration by Nmax (Mmax MTB (etB)//Mpt MM (see Table 2 In the attached document, we provide this table with tools to alter values (see Appendix B.2.1
We configure the axis sizes Table improve performance. Lower axis sizes allow for more warpgroups. This increases the degree of memory sharing, akin to increasing J7 in Figure 44 However, with g7 2 295, not bandwidth bottlenecked H100 (see Appendix B.0.1 even assuming caching. Larger tiles
16

---

Published in Transactions on Machine Learning Research (03/2025)
Variable
Size
Q.
Level
MTB (Byles)
MwG
(Byles)
Q
Queries
w 128) x
FP8
SMEM
16384
K
Keys
(32) xf(228)
FP8
SMEM
16384*
V
Values
(32) (128
FP16
SMEM
32768*
S
Attention Scores
wI(288) (32)
7 X
FP16
Registers
16384
P
Weighted Scores
97
(228)
Sqz
FP16
SMEM
16384
P'
Weighted Scores (Registers)
by
FP16
Registers
16384
A
Weighted Scores (Tensor Core)
Xux "
(128)
FP16
SMEM
8192
a
Auxiliary
4x
FP16
Registers
768
Output (Registers)
txxd
FP16
Registers
32768
D
Delta Output (Tensor Core)
(1128) xx6
FP16
Registers
8192
A0
Delta Output (SMEM)
(28) d((16)
FP16
SMEM
8192
AO'
Delta Output (Registers)
1xxd"
FP16
Registers
2048
Total (KB=1024B)
SMEM
48
48
Total (KB)
Registers
74.75
Table 1: The configuration table lists all the required variables and the memory they occupy. We set why 128, tq = 8x Ug 64, d = 128, d' 32, and d" 8. Alternative configurations can be explored in the linked document. *Keys and values are asynchronously loaded, doubling their memory usage.
allow for greater parallelism and simplify the code, motivating the compiler to cooperate. The number of non-tensor FP16 multiply-add operations is dependent on the number of sub-loops, which we reduce by setting uy = $g. If we were working with FP16 queries and keys, we would have aim for more memory conservative configuration, possible docrcasing the size of subloops. Configuration tables allow to explore these possibilities in a systematic manner. In the associated document (Sheet "Configuration"), we include file that allows configurations to be quickly assessed.
Excess at N =
M Miax
MTB
MwG
N
Per TB.
Per WG.
Per thread (Byles)
SMEM
227
48
48
3.7
35
11.67
Registers
256
0
74.75
3.4
31.75
10.58
85
Table 2: We find the maximum number of warpgroups given our configuration using a simple linear expres- sion, Nmax (Maax MIB)/MwG. The excess registers per thread provides buffer to prevent register spills which devastate performance.
5.7 Throughput Optimization
We can use diagrams to identify the bottlenecks of algorithms to guide optimized overlapping strategies. Each SM has specialized cores and [pipelines for different operations which execute in parallel each clock cycle. GPUs will naturally implement a degree of overlapped execution, however, FlashAttention-33 (Shah et al. 2024 showed that explicit barriers can improve asynchrony.
From Figure| 29 we find the ops per thread for each column. We take the baseline ops (eg. 2d for matrix multiplication) and multiply by the weaved axes. We ignore Φ7. to normalize for the number of active threads. In Table 3 we notate the pipeline each operation uses, the normalized ops, and the ops per clock cycle. We use these values to find the required clock cycle per column. These values are given per & subloop. The compute analysis is provided as Table 3 in the altached document. for various algorithms (see Appendix B.2.1..
Table indicates that clock cycles per iteration are lower bound by the 6 clock cycles per thread required for tensor cores. Waiting any other actions will deviate performance from this lower bound. We can
17

---

Published in Transactions on Machine Learning Research (03/2025)
Operation
Q-K MatMul
SoftMax (erpo- nenl)
P-V MatMul
FP16 Accumu-
late
Pipeline
Tensor
Special Func- tion Unit
Tensor
FP16
Ops/Th
2ds; 16384
$5 + 55/ 65
28dd 16384
2d ST U7)
256
Ops/Clk 8192
16
4096
512
Cik/Th
2
4.06
0.5
Table 3:
We outline the clock cycles per thread required to execute different stages of the algorithm. The total clock cycles required are given by (CIk/h))(Nuubbll of Active Threads).
this lower bound to find the required warp groups to avoid being bandwidth limited. Assuming no caching, we find that 97 must be greater than 295. This is claborated in Appendix B.0.1 where we also find the ideal tensor-core bottlenecked throughput to be 1.32 PFLOPs.
Hopper provides synchronization mechanisms between warp groups, which we exploit to shadow non-tensor core operations. Optimal performance is achieved by having barriers wait on tensor cores. We diagram a representation of the algorithm in Figure 34 with widths corresponding required clock cycles per thread. We classify costs not considered in Table slow] and exponential operations must occur in FP32. We represent overhead accommodation with hatched regions.
overhead. For example, [ype conversions extremely
Figure 34: We can diagram the clock cycles of Figure 33 with a specific configuration (uz 8F, d' = d/2) by symbolic blocks, with width proportional to the clock cycles per thread. We use d' d/4, but the strategy remains similar. Hatched areas indicate overhead, which we assume to be significant for non-tensor core operations. The ruler width is equal to the Q-K FP8 tensor core matrix multiplication, and therefore represents 1 clock cycle per active thread.
We expand the diagram into three warpgroups. Each warpgroup acts independently, not relying on data from others, allowing us to shift operations forwards or backwards to wait on tensor cores. Furthermore, can split the Q-K tensor core matrix multiplication allows us to construct highly overlapped algorithm with two active iterations, shown in Figure 35 This
it forms linear subloop (see Figure 30
This
algorithm will achieve optimal performance if the realized overhead is less than the accommodation sizes.
DIDI
D
D
Figure 35: Three warpgroup pipelining strategy. Hatched areas indicate accommodation for overhead of non-tensor core operations. We have 50% overhead for Soft and 100% for FP16 multiply-add. The ruler separates blocks of 128 clock cycles. Thick dotted lines indicate barriers. the ideal case, these should all wait on tensor cores. We have two itcrations active once.
18

---

Published in Transactions on Machine Learning Research (03/2025)
5.8 Conclusion of Derivation
In this section, we introduced systematic method for deriving, configuring, and analyzing algorithms for our toy Hopper hierarchy. The utility of this approach is in quickly deriving sketches of optimized algorithms with clear idca of expected performance and how the algorithm can fall short. The diagrammatic methodology can be extended to simulate additional architectures for which kernels have not yet been developed, including new releases like Blackwell [NVIDIA 2025) and non-NVIDIA hardware (AMD) 2023
The toy model, however, overlooks certain features of GPU algorithms. Though diagrams provide accom- modation for overhead, we have ignored the overhead of tensor core operations and assumed that latency is completely hidden. In practice, latency may delay algorithm exccution, and tensor core operations may incur overhead from small tile sizes (Bikshandi & Shah 2023). The aim of this work is to develop a theoretical framework that allows specific assumptions (like tensor core overhead) to be improved by empirical testing, Rather than generally stating "larger tensor core tile sizes improve performance", we may be able to claim that "the overhoad of tensor core operations is Y given a tile size X
A major concern is our lack of utilization of tensor crageeengee register-level operations. We have modeled tensor core memory as incoherent, however, methods exist to find how fragments are stored across registers, which can be exploited to avoid intermediate communications with SMEM (Bikshandi et al. 2023 This requires special treatment, as we would have to consider reindezing operalions which manipulate how data is accessed without changing values. This can be encompassed by a reindexing weaves (Abbott) 2023
which can capture the implementation details of CUTLASS layouts and swizzled memory (Shah 2024
6 Comparison to FlashAttention
Our diagrammatic approach and "ideal" algorithm provide insight into FlashAttention-8.. The key difference in the operations employed is that FlashAttention-33 manipulates tensor core data thread-wise. This allows it to use less variables than our approach, as we document in Appendix B.2 This functionality is not currently provided by diagrams, meaning our algorithm resorts using SMEM data. Properly considering the manipulation of tensor core data would require investigating indexing with diagrams in-depth. An ad- hoc treatment would detract from the emphasis on systematie methods of this work. However, the value of diagrams lies in the systematic methodology and analytical tools, which we can be applied to better understand FlashAttention.
6.1 FP16 FlashAttention-3
FP16 FlashAttention-3 uses FP32 auxiliary variables and intra-warp group overlapping. We outline its configuration in Appendix B.2.4 linking to the attached document. Large tile sizes are used, with effective 97 = 128 and $7 128. Intra-warp group overlapping requires single warpgroup to hold gnext, overloading the limited number of per register threads, causing spillage
We diagram the intra-warpgroup strategy and a hypothetical alternative inter-warpgroup approach for d == 128 in Figure 36 We see that the inrra-warpgroup approach does not accommodate overhead for SoftMax, potentially delaying the algorithm. The inter-warpgroug strategy provides 100% overhead for MaxtMax. This invites empirical testing where we roduce the tile size to fit multiple warpgroups per SM. If an intra- warpgroup strategy remains superior, we can update the assumptions of our toy hierarchy based on the empirical results.
The choice of large 8p is the result of empirical testing and trado-off between register pressure and the performance benefits of large tiles. This indicates our model can be improved by considering the overhead of tensor core operations, especially for small tile sizes. Without caching, the algorithm will be bandwidthe bound. Nonetheless, the algorithm is performant, achieving 740 of the maximum 989 TFLOPs of compute (75%). Diagrams like Figure 32 show quantization changes more clearly than code, and we note that it is likely the case that much of the value of FP32 accumulators is lost in the second matrix multiplication.
1 Exxeeeding register limits causes local memory to be used. Despite its name, is located off the SM
19

---

Published in Transactions on Machine Learning Research (03/2025)
FP16 Inter- Warpgroup Overlapping
FP16 Intra- Warpgroup Overlapping
FP16 inter and intra warpgroup overlapping strategies. Hatched regions indicate SoftMax overhead.
Figure 36:
6.2 FP8 FlashAttention-3
FP8 FlashAttention-3 uses FP16 auxiliary variables and inter-warpgroup overlapping in the original version, documented in Appendix B.2.2 However, newer versions seem to use the intra-warpgroup overlapping operation (see Appendix B.2.3). This strategy allows a single warpgroup to store a large 8z value, which seems to have an immense benefit. FPS FlashAttention--3 bound by the slow special function unit required for the SoftMax function. empirical testing, achieves 1.2 PFLOPs of the maximum 1.98 (60% utilization).
In Figure 37 diagram the overlapping strategies for FP8. We overlapping operations and assess the impact of function overhead. We provide diagrams for both inter- warpgroup strategy (original version) and an intra-warpgroup approach (mimicking the GitHub). The clock cycle analysis shows that the strategies are both delayed by the overhead of the SoftMax operation. As the SoftMax operation utilizes as many clock cycles as tensor core operations, this can not be improved upon. Delays from unit couversions and other factors will delay the algorithm, reducing the maximum achievable FLOPs. Furthermore, the FPS intra-warper warpgroup strategy has the special function unit (the limiting factor) inactive for much of the cycle, reducing performance by minimum of 33% which explains much of the underutilization.
clock cycle assessments to show
FP8 Inter- Warpgroup Overlapping
FP8 Intra- Warpgroup Overlapping
Figure 37: FP8 is bottlenecked by the SoftMax operation, and therefore has no accommodation for its overhead, no matter the overlapping strategy employed.
20

---

Published in Transactions on Machine Learning Research (03/2025)
6.3 Further Comments
The aim of this analysis is to see how our assumptions and predictions can be compared to a prior successful kernel. The choices made for F especially the latest releases, of warpa-wap strategies, large SI tiles, and a small number of warpgroups 97, suggests that there is a significant benefit to large inner dimensions for lensor core operations. If this is the case, our model would need to be updated. This does not detract from this work's main goal we aim to propose theoretical framework that can be updated and tested in exactly this manner. In the long term, the exact details are secondary to the mechanism that allows them to be updated. By positing certain assumptions about constructing an optimal algorithm as in Section and seeing how it falls short, we can inform future development.
Our model predicts that FlashAttention-3 bottlenecked by the SoftMax operation. Interestingly, if we assume a 66% overhead in SoftMax, we would expect FP16 intra- warpgroup to take 4/3 as long as necessary and FP8 inter warpgroup take 5/3 as long as necessary, explaining both the 75% utilization for FP16 and the 60% utilization for FP8 in the original paper. This is alternative hypothesis from delays originating from the overhead of small tile sizes, and invites empirical testing.
7 Conclusion and Future Work
In this work, we have contributed a diagrammatic method for representing deep learning algorithms. Tn Section [ we showed how this allows for the rapid derivation of high-level optimization strategies, with an in-depth associated performance model shown in Soction 4 In Soction 5 we used this framcwork to develop systematic hochoogyg for deriving hardware-aware algorithm. In Section F we use our methodology to better understand FlashAttentonn, making testable claims about the factors that limit performance.
This work is a zero-to-one attempt at creating systematic approach to understanding IO-Awareness, and invites a scientific approach to GPU optimizations. Diagrams rely on assumptions about GPU behaviour to predict real-world performance. By separating theory from empirical results, we are able to treat these assumptions as scientific hypotheses. Experiments can be run, with their failure or success guiding im- provements in the framework. The diagrammatic framework can therefore avoid post-hoc rationalization. Furthermore, diagrams reveal the many possible configurations of algorithms in Figure 30 and 32 These configurations can act as templates over which host of experiments can be run. As the framework becomes increasingly capable of predicting performance, the cost of empirically testing various configurations can decrease.
This work focuses on attention. However, diagrams can provide value by extending to other algorithms. Mixture-of-experl models (Mistral AT leam 2024) use immense resources, making optimizations particularly impactful. Additional strategies can be formalized. Convolution and sliding window attention (Beltagy al, 2020) reindex weavings (Abbott] 2023). operations which change how data is accessed without changing the data itself. Reindex weavings can integrate CUTLASS tensor storage system linking diagrams to CUTLASS implementations and providing access to tensor core-level weaved operations.
Diagrams conform to a category-theorelie description, which is not covered in this paper. A categorical perspective would have diagrams encoded as a formal syntax (Piedeleu & Zanasi 2025) which opens up automated compilation tools (Wilson) 2023), ; consideration of backpropagation (Fong et al. 2019 Cruttwell et al.] 2022), and estimating error accumulation (Perrone) 2024). A formal treatment would streamline our theorems, and provide access to tools that allow additional strategies to be considered. Convolution and sliding window attention (Beltagy et al. 2020) use reindexings on weavings. These operations change how data is accessed without changing the data itsclf, and correspond to Yonoda natural transformations from category theory (Abbolt, 2023). Reindexings would allow CUTLASS' tensor storage system to be incorporated into diagrams, and introduce a systematic method for optimizing convolution-lis algorithms. Furthermore, we can develop methods to distribute streams across cores and use accumulators as multi-core reduction operations (Osama al. (023). ceeumllaors-as--rcdetiii requires special treatment of transfer costs, but would allow wave quantization to be considered which is outside the scope of this paper.
21

---

http:p/rreeeerrrr...ooooo/oieocs Published in Transactions on Machine Learning Research (03/2025)
Furthermore, can use our framework to develop systematic approach to optimization that considers hardware design, capabilities, and algorithm configuration to performance. This can be constructed by adapting the framework of categorical co-design (Zardini 2023), which relates functionalities to requirements in a compositional manner. Given the categorical nature of diagrams, this would the development of a holistic approach across the engineering stack of opt imizing deep learning deployment, relating hardware, algorithms, and configuration choices towards optimal performance and minimal resource usage.
References
Vincent Abbott. Robust Diagrams for Deep Learning Architectures: Applications and Theory. Honours thesis, The Australian National University, 2023. URL https:/ http://wwww.abbbblloooooonnnt /f0p:///een///e/2222///// 1
Robust- Robust-Diagrann Diagrams- Deep- Learning Archineecciiteetre.p.d.ff
Alok Aggarwal and Jeffrey Vitter, S. The input complexity of sorting and related problems. Commun.
ACM, (1(9):1116 (9111111111 1988. ISSN 0001-0782. doi: 10.1145/48529.4800 URL butpss//ddlacms /oi/ 10.1145/48529, 44.144544559448533
Joshua Ainslie, James Lee-Thorp: Michiel de Jong, et al. GQA: Training Generalized Multi-Query Trans- former Models from Multi-Head Checkpoints. In The 2023 Conference on Empirical Methods in Natural
Language Processing, 2023. URL https://openreview: net foot//foummiiii fowummidd-mmwwwwWWW
AMD. AMD CDNA 3 architecture, 2023. URL httppp//wwww..aan.mo////otttt/// http:///aaaaaaaama/am///cccoaaae instinct-tech-doos white- wwwte-papers////eeeee papers/amd-cdnnse 3- white- paperr-dp
Iz Beltagy, Matthew E. Peters, and Arman Cohan. Longformer: The Long- Dong-Document Transformer. arXiv preprint, 2020. URL https:/ /axtss http:/iwwwiiar.aaaaa//eea wowgoorg//////o..oo.o..oo. habs/2004..05111
Ganesh Bikshandi and Jay Shah. Developing CUDA Kernels for Accelcrated Matrix Multiplication on
NVIDIA Hopper Architecture using the CUTLASS Library. Colfax Research, 2023. URL htLpst/s::/// httpp://ressarrhh research. gemm-kernels- menes-intl..com//eee www.:o///ttwootott//oooommaaoeee wpwpadpp /2033/ 12, colfax- wowfax-eeeaa----- gemm- kernels- kernels--oppere hopper.ppl
Ganesh Bikshandi, Jay Shah, and Colfax Research. Case Study in CUDA Kernel Fusion: Implementing F on NVIDIA Hopper Architecture using the CUTLASS Library. Colfar Research, 2023.
2203 12 /aaa//////llll/aaaaaaattntnnnnno
Geoffrey S. H. Cruttwell, Bruno Gavranović, Neil Ghani, et al. Categorical Foundations of Gradient- Based Learning. In Ilya Sergey (ed.), Programming Languages and Systems, pp. 28, Cham, 2022. Springer International Publishing. ISBN 978-3-030-99336-8. URL https:: https://linnk.srinnger htt.e/ee.cc//////.//eeeeee coom/chapper 10. 1007 00779978-30000-99330
Tri Dao. FlashAttention Faster Attention with Better Parallelism and Work Partitioning, In The Twelfth International Conference on Learnång Representations, 2023. URL https://oopnnreview.nttt forumm/ummm ?id-
mZZ22yh9Ec
Tri Dao, Daniel Y Fu, Stefano Ermon, et al. FlashAttention: Fast and Memory-Efficient Exact Attention with 10- AO-Awneessss. In Alice H. Oh, Alekh Agarwal, Danielle Belgrave, and Kyunghyun Cho (eds.), A Advances in Neural Information Processing Systems, 2022. URL
net Htae:/oommmdddd--DDDaPpa-Siaaa
Patrick Esser, Sumith Kulal, Andreas Blattmann, et al. Scaling rectified flow transformers for high-resolution image synthosis. In Proceedings of the 41st International Conference on Machine Learning, ICML/24 244. JMLR.org, 2024. URL htlp:://ddlacmmorr..
Brendan Fong, David I. Spivak, and Rémy Tuyéras. Backprop as Functor: A compositional perspective on supervised learning. In 94th Annual ACM /IEEF Symposium on Logic in Computer Science, LICS 2019, Vancouver, BC, Canada, June 24-27, 2019, pp. 1-13. IEEE, 2019. doi: 001108//LICS.2011.88558888 URL htlps://dlaam..o//ddo////
22

---

Published in Transactions on Machine Learning Research (03/2025)
Elias Frantar, Salch Ashkboos Torsten Hoefter, and Dan Alistarh. OPTQ: erative Pre-trained Transformers. In The Fleventh International Conference on Learning Represenlations. 2023. URL [ttpps//ppeeneeiee.eeef//oo forummiid=ttiBBPiiis..
Accurate Quantization for Gen-
Amir Gholami, Sehoon Kim, Zhen Dong, al. A Survey of Quantization Methods for Efficient Neural Network Inference. In Low-power computer vision, pp. 291-326. Chapman and Hall 2022.
Amir Gholami, Zhewei Yao, Sehoon Kim, et al. AI and Memory Wall IEEE Micro, 44(03):33 39, May 2024. ISSN 1987-4143. doi: M0009/MM..224.3773399 WM0/MMM.2024.377777
URL https: https:://dii.cceeeee http:://ooo.ceccommttttttttttttt 10.
Saugata Ghose, Abdullah Giray Yaglikgi, Raghav Gupla, et al. What Your DRAM Power Models Are Not Telling You: Lessons from a Detailed Experimental Study. Proc. ACM Meas, Anal. Comput. Syst., 2(3), December 2018. doi: 10.1145/3224419 URL https:/ /doioo////////////////////o 10.
Jonathan Ho, Ajay Jain, and Pieter Abbeel. Denoising Diffusion Probabilistic Models. In Hugo Larochelle, Marc' Aurelio Ranzato, Raia Hadsell, et al. (eds.), Aduances in Neural Information Pro- cessing Systems 39: Annual Conference on Neural Information Processing Systems 2020, NewrIPS 2020, December 6-12, 2020, virtual, 2020. URL https:// httpss//proceedingsneeerrrp
/20pr/2020///// hash/
79ca4b- Abstractt.tml
Andrei Ivanov, Nikoli Dryden, Tal Ben-Nun, et al. Data Movement Is All You Need: A Case Study on Optimizing Transformers, 2021. URL http:/ www.aviwvvoog/...aa httpg//abb/200..0/000 2007.00072
Llama team. The Llama 3 Herd of Models, 2024. URL http: aatpaaaaiiii.//./aa//////// attaaabb/242022.2 2407.21788
Weile Luo, Ruibo Fan, Zeyu Li, et al. Benchmarking and Dissecting the Nvidia Hopper GPU Architecture In 2024 IEFE International Parallel and Distributed Processing Symposium (TPDPS), pp. 656-667, Los Alamitos, CA, USA, May 2024. IEEE Computer Society. doi: 10.100//PDDPS577555222....0 URL
Mistral AI team. Mixtral of experts, 2024. URL attp:s//mustrrlaa//ee////tttktos
NVIDIA. NVIDIA A100 Tensor Core GPU Architecture Overview, 2020. URL https: docs..vidiaa.cmm
liippptxpp__po pff
NVIDIA. NVIDIA H100 Tensor Core GPU Architecture Overview, 2022. CRL hitl https//rrsouurces.nvidaaa
com com/en- eeeeentt tae--us-eeee tensor-core
NVIDIA. PTX ISA 8.5, 2024. URL
isa 8.5.pdf
NVIDIA. NVIDIA Blackwell Architecture Technical Overview, 2025. URL
en en-us-blacken hwt.dd/www..eiiiiie////////iaaae
Hiroyuki Ootomo and Rio Yokota. Reducing shared memory footprint to leverage high throughput. on
Tensor Cores and its flexible API extension library. In Proceedings of the International Conference on
High Performance Computing in Asia-Pacific Region, HPCAsia '23, pp. 8, New York, NY, USA, 2023. Association for Computing Machinery. ISBN 9781450398053. doi: 111445/577118.3558222 URL https:
//doi.org/ 1oniiorg/10.11140 35114555/788835577788
Muhammad Osama, Duane Merrill, Cris Cecka, et al. Stream-K: Work-Centric Parallel Decomposition for Dense Matrix- Matrix Multiplication on the GPU. In Proceedings of the 28th ACM SIGPLAN Annual Symposium on Principles and Praclice of Parallel Programming, PPoPP '23, pp. 429 431, New York, NY, USA, 2023. Association for Computing Machinery. ISBN 9798400700156. doi: 10.145//572848.3577777 URL https://do..rrg11.114///////
Matrix-Matrix
Adam Paszke, Sam Gross, Francisco Massa, et al. PyTorch: an imporative style, high-performance deep learning library. In Proceedings of the 33rd International Conference on Neural Information Processing Systems, Red Hook, NY, USA, 2019. Curran Associates Inc. URL https://pytorchorrgg
23

---

Published in Transactions on Machine Learning Research (03/2025)
Paolo Perrone. Markov Categories and Entropy. IEEE Trans. Inf. Theor. 70(3):1671-1692 March 2024.
ISSN 0018-9448. doi: 10.1109/TTI /01109/TIT.OO 33228332882.. URL hitpss// /doi.o////////////// 10.oo.oor/1011100 10. TT0T100//TT.22223333258888 2023.3328825
Robin Piedeleu and Fabio Zanasi. An Introduction to String Diagrams for Computer Scientists. Flements in Applied Category Theory. Cambridge University Press, 2025. URL lttp:: http:://ariix.org/as////50
Dustin Podell, Zion English, Kyle Lacey, ct al. SDXL: Improving Latent Diffusion Models for High-Resolutee Image Synthesis. In The Twelfth International Conference on Learning Represenlations, 2024. URL
net foorumidd--a di12:ddddddd22xxttt
Robin Rombach, Andreas Blattmann, Dominik Lorenz, et al. High-Resolution Image Synthesis with Latent Diffusion Models. In 2022 IEEE/CT /CVF Conference on Computer Vision and Pattern Recognation (CVPR), pp. 10674-10685, Los Alamitos, CA, USA, June 2022. IEEE Computer Society. doi: 10.1109 10.1109/CVPPE CCVPR2688.. 10. /t/////VV22..2222....22
2022.01042. URL
Amit Sabne. XLA Compiling Machine Learning for Peak Performance. Google Res, 2020.
URL http:
Jay Shah. A note on the algebra of CuTe layouts. Colfax Research, 2024. URL https: htipss//roscarcc.oollaaa intl.
com/ prp-coonlen/uuploeee 2024/01 /layout algebra.pdf
Jay Shah, Ganesh Bikshandi, Ying Zhang, et al. FlashAttention-3: Fast and Accurate Attention with Asynchrony and Low-precision. In The Thirty-eighth Annual Conference on Neural Information Processing Systems, 2024. URL https://openreview..t net http://wwww.ee/ooonniiiiiee /tt////rottt///CC Yid20
Noam Shazeer. Fast Transformer Decoding: One Write-Head is All You Need. arXiv preprint,
(arXIv:1911.02150), 2019. doi: 10.48550/ 00048850/arXiv.1111.. 191iv.1911.02150 URL http:/ httpp///axiv.or..gp aaaaaabss11111 1911.02250
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, pp. 10-19. Association for Computing Machinery, URL https:/ https////laaac.oo/ooo/000.0.0
2019. ISBN 978-1-4503-6719-6 doi: 0114553311508.33222 3315508 3320978
Ashish Vaswani, Noam Shazeer, Niki Parmar, et al. Attention is All you Need. In Tsabelle Guyon, Ulrike von Luxburg, Samy Bengio, et al. (eds.), Advances in Newral Information Processing Sys- tems 30: Annual Conference on Newral Information Processing Systems 2017, December 4-9, 2017, Long Beach, CA, USA, pp. 5998 6008, 2017. URL. https:// httpss//poocellnggnnerrrig paper /211 (017//10 hash/
Abstraac.ttt
Paul William Wilson. Category-theoretic data structures and algorithms for learning polynomial circuits. phdthesis, Universily of Southampion, 2023. URL http:://eepriilssooooo.....k../.
Gioele Zardini. Co-Design of Complex Systems: From Autonomy to Future Mobility Systems, 2023. URL Accepted: 2023-12 19700:3:577.
https:/ wwtp://searrdd-ooolltoooo..eiiia handle 1tw.cle/2..000.1155//66
A Appendix
A.1 Fusion Theorem
Theorem (Fusion Theorem) Composition and weaving of a streamable algorithm which does not re- move the streamed aris yields streamable algorithm.
Streamable algorithms satisfics the form in Figure 38 allowing the recursive expansion of Figure 39 Stream- ability depends on polymorphism over the a axis, meaning the function function/algoritaa algorithm is defined for the axis being of any size, and the existence of an accumulator B, which allows the streamed input data to be split.
24

---

Published in Transactions on Machine Learning Research (03/2025)
Figure 38: A streamable algorithm quires an accumulator B which allows the polymorphic streamable axis to be split. We can fuse the algorithm with a head and tail, which do not require ad- ditional loads and saves if their memory usage sufficiently small.
b
A
b
b
a
F
x
a
F
8
B
X
This allows for recursive decomposition, reducing the size of the remainder a' axis to a' S 8a as Figure 39 Here, we add a head and a tail, which are not expanded but can be executed without an additional transfer and are hence fused. This requires their memory usage to be sufficiently small limitations. Composition on G K simply replaces those algorithms with the composed form, yielding modified head/tail for the streamable algorithm.
exceed hardware
b
b
a
2s,
Sg
Sq
Toil. @
F
K
Head
a''
F
y
Sq
B
y
y
B V
K
V
G
x
Head, a
Figure 39: Figure 38 is satisfied, then the algorithm can be recursively decomposed. This reduces the size of chunks until the on-chip memory is sufficiently small.
Composing by E' on the b-axis with an algorithm weaved by the streamable sa axis, we can exploit a partition copy (see Figure 10 to show that the composed F' has an accumulator B'.
b
E
a
S
b'
E
Thil,
w
b
F
M
K
y'
Head
a'
F
S
B
Toil
G
G
K
Head, a
b
E
a
S
S2
g
E
9
Tail
b
S,
F
K
V'
Head
E
F
S
B
Tail
G
K
Head, a
b'
b'
a
Sp
B
b'
b
Tril
E
b
b'
F
K
V'
Head
F
E
B
Tail
x
G
G
3
K
V
Head, a
Finally, we are required to show that weaving preserves streamability. This exploits characteristic of map- ping composed functions. Mapping composed function additional axis is equivalent to composing
25

---

Published in Transactions on Machine Learning Research (03/2025)
the individual functions mapped over that axis. Therefore, show that weaved B is an accumulator for weaved F as in Figure 40
Tail
b'
b
b'
E
q
S
B'
F
K
Toil
x'
G
E
b
E
7
Head
F
S2
Head
G
B
Y
K
V g
Figure 40: Weaving over F composed with E and G at its inputs is equivalent to weaving over the smaller F and B algorithms which, when composed, give F. This weaves need to trace over the inputs which they target.
We can combine all the above expressions into the single form of Figure where we also apply group partitioning. This separates the mapped axis into groups of size distributed across processors. We can iteratively apply these rules; an algorithm can be composed with an algorithm E weaved over the streamable axis. This generates a strcamable algorithm, which can be weaved over one of the inputs introduced by E. This is used to construct streamable (lash) atlention from SoflMaxx-Cnnraatioon kernel.
8
8
8
Taxl
6'
0'
b²
b
E
B'
Sn
5
F
M
K
y
Tail
b
b
x
E
gr
Head
F
B
Ilead
(gg1g,)
(
G
Y
K
g
X
Figure 41: We can combine the streaming theorems into a single expression. Given that the algorithm is streamable along the a/sa axis, add modifications to generate streamable algorithm. We group partition along the streamed axes, mapping groups of size 9q gr to different processors.
A.2 Multi-Level Performance Models
We define for two-level model diagram the optimal transfers (a M), where ā are the axis sizes and M is the maximum memory available at the lower level. We use ÿ to indicate some configuration of group sizes. We are interested in memory usage in the limit of large M and J, in which case the limiting factor is some y П. 9: weaved by all grouped axes, with y typically being the output. For smaller M, we aim for a specific configuration as in Section 5 Two-level diagrams derived optimal transfers solve for Equation where iterates over the grouped axes:
The as II gi Hz given M 2 y II, 9:
Ha min g
(4)
We extend this multiple levels be assigning memory Mt each level and a weighted transfer cost H. which represents bandwidth. As described in Section this conforms to;
26

---

Assuming that the inequalities are Published in Transactions on Machine Learning Research (03/2025)
H' =² = H (a,
(5)
How additional levels are used, cither as intermediate caches or cross-transfer levels, use the same general performance model of Equation 5 but with altered weighted transfer costs H and per-level effective memory My. Using M to indicate the maximun memory per level, N to indicate the maximum number of child nodes per higher level node, h for the immediately higher level and for the immediately lower, and Hige as the raw weighted transfer cost between l and f' we adapt our performance model to get;
1 Hele 1
is a cross-transfer level else
H
Nunax Me
l is a cross-transie or intermediate cache
Mr
M
else
Therefore, using cross-transler or caching levels conforms to our standard performance models. This lets Equation act as a universal performance model for multi-level GPUs, and ensures that H M, is the characteristic property for comparing the performance of algorithms for different low-level memory sizes.
A.2.1 Intermediate Caching
Theorem An intermenliate caching level E for oulput restricted algorithm with number restric- tion Niz N9,22/Ng.a conforms to the standard performance model with M Me2NNux requiring H' (ā. N (22x)) lotal transfers.
For an intermediate caching model which is output limited, we have the standard constraint for the lower level, Mr2 > 11, 9e2,,, and the number restriction, N > Ng,22/N,, = 11, ge1,//l,, 922,i- We aim to find the configuration of Ha and gez which minimizes the number of transferred values. Assuming the algorithm is output limited, the effective size of the caching column y IT, gf This lets us express the restrictions as:
II, 0: II ga.
HA min
Hz(ā,ge) fe)
given Nmax > II 9 gea.i
(6)
II, 02 (d
H = min
given MA2 2 II 922,5
(7)
II, gez,i
We can multiply the restriction of 6 by the restriction of 7 get sufficiently close to equalities, outline the problem to solve:
II. a: (a ga1) 11, gel, I.: ai
H min
given Ma > II; GEl.s
(8)
H min
Hy (ā. 922)
given M1 y II gf²
(9)
II; g82,i
These restrictions correspond to Equation 4 Nmax for the intermediate level instead of its own maximum memory, Mmax We therefore have;
substitute in H Me) for both levels, but using
H H Mr2
H = (a, Mez)
A caching level therefore conforms to our standard multi-level performance model derived from two-level diagram, but with the intermediate level using total lower level memory instead of its own.
27

---

Published in Transactions on Machine Learning Research (03/2025)
A.2.2 Cross- Transfer Level
Theorem Introducing a cross-transfer level a between a higher level and a lower level e allows us to replace the weighted transfer cost of an output-limited algorithmm;
H = - (a, + ...
With,
H'=.. Where Hn is the weighted transfer cost of sending data to children, and Hybe is the weighted transfer cost of sending data between children.
(a, M.) + Ha
The child level requires total of H¹ Mc) data transfers from the higher level, typically incurring a weighted transfer cost per value of Hn With a cross-transfes level z between h and we can send some data Hº to any of the children and cross-transfer the remaining at a weighted transfer cost of H per value. We need to derive the optimal configuration of Te to minimize H² This configuration incurs a number restriction, the child processors need to remain active. We therefore have,
II; as Tx) 11, Jo,
H = min
given N > IIs Iz.i/ II, 9c.
d: (ā.g.)
3 = min
given Me 7 11; 9c,i
# go.i Go.i
We perform similar substitution to Section A.2.1
yielding new of restrictions;
II, 0j Hd gx)
H = min
given M.NYTaxx > II Iz.i
II, 9x.
11, a IT. 9o.
H. min J
ga)
given M. R y II; Sos
These restrictions conform the two-level model optimal, with required transfers being H (ā McN and Hº (ā, Mc). When considering transfers for the total weighted transfer cost caleulation, we subtract the required transfers to from the transfers required to This is because data is transferred from the higher level to the cross-transfer level by sending it to children, so much of the data is already available. This results in the substituting the total weighted transfer costs 10 with
H' =... + 1 H' (ā, Mo) + 11 .... +1 N Mo) + 11 (H" (,, Mc)) +
(10)
(11)
We can rearrange 11 80 that we have the standard format of each level having H' (α, Mr) transfers by instead modifying the transfer cost:
H = + (H - H (ā, N.aa M.) + #-1 2-t0 (T, +
(12)
Therefore, a cross-transfer level conforms to the standard performance model. Instead of using the weighted transfer cost of sending data to the children 11,11 for the cross-transie level, we set 11 H 1=
and we remap H. H This lets us express 12 as the expression below, which conforms to the standard multi-level performance model of 5
= HJ! H Ha Nmax Mc) H-1 Ha
A.2.3 Additional Notes
In the case of a multi-GPU hicrarchy, the intereonnected topology is a cross-transfer level which distributes data among child GPUs at weighted transfer cost of H If we assume data already distributed across GPUs, then the number of GPU cross-transfers is H (α, Mflax) H* Nmax Mmnx). So far, we
28

---

Published in Transactions on Machine Learning Research (03/2025)
have considered the highest level to have unlimited memory and zero weighted transfer cost. We can model multi-GPU systems by the highest level (the multi-GPU level æ) as having memory equal to Nemax Ma and negative weighted transfer cost -Ht which assuines data is already distributed among GPUs. This provides a rough model, which can be refined by aligning the group sizes used in diagrams.
Often, we can configure the number of cross-transler level children alter the cross-transfer bandwidth. Tn (uu al. 2024), was found that the bandwidth of H800 ( (Chinese market Hopper GPUs) SM-SM transfers varies from 3/27TB/s s with cluster size N = 2 to 2/ with cluster size of N = 4, compared to B/ of GMEM-SMEN bandwidth. This imposes a trade-off: smaller cluster sizes improve effective He but reduce the cross-transfer discount H" Nimax NMtasMMaxx), (Luo et al. 2024g notes that balancing this trade-off is an important direction for exploration. We can use our model to find the difference in weighted transfer costs with 11 and without [10 a cross- transfrr level, providing an equation to optimize for N.
AH' =
(á, Mc) + ...)
- -(... + -... + - -H H (ā,
Mc) + - Ha( ...)
(H - -Hpbo) H Mo) - -H - -H (a, Mo)
-1itihe- - -Httd) (H (ā,, Ma) - HH N Mc)
AHI (N) Σ al(()) M.0 0 (1-N-9)
A.3
Streamability
A.3.1
Contraction
The streamability of contraction (a dot product) requires that an accumulator exists of the form in Figure
38 Contraction for vectors v, w E Ra is given by v- w = 2/- w., which can be expressed as w
D 02 w; + ---1 0/ Wi, where the underlined portion is the accumulator. We diagrammatically show this in Figure 42 highlighting the accumulator in blue
Figure 42: We can re-express contraction as an initial function followed by the accumulator. Any size can be chosen for and s', and the expression can be recursively expanded until 5¹ S 8, for some target stream batch size 8a.
a
A.3.2 Fusion of Contrax-Contratton
SoftMax-Contrration is strcamable by maintaining a running maximum and sum on chip as auxiliary vari- ables. We express this by streaming unscaled SoftMax- Conttraconto with the initialization of the auxiliary variables as the head and scaling by the sum tail as in Figure 43
Streamable SoftMax- Contraction is implemented by accumu- lating the results of unscaled SoftMax- Contraction applied to segments, adjusting the baseline relative to the current maxi- mum value. We apply the normalization by as a tail for the expression.
Figure 43:
X
X
1
29

---

Published in Transactions on Machine Learning Research (03/2025)
We can derive streamable SoftMax-Contrachi by taking recursively expanded Auxiliary SoftMax (Figure 44h and contracting its output. Recursively expanded Auxiliary SoftMax is not streamable the memory usage increases with the imput size. However, it terminates in join, allowing it to be fusee with contraction.
Auxiliary SoftMax (defined in Table where we maintain auxiliary vari- ables, can be recursively expanded.
Figure 44:
5
2012
3
3
3
Base SoftMax
Auxiliary Soft
Unscaled SoftMax
x
x
x
x
X
Max))
Initialize()
Initialize()
" f max max)
Return (-00, 0,0)
Return (-00, 0, 0)
8; exp( (zz; 1t)
SoftMax, (p': 8z', 2²)) })
SoftMax, (5, (p',8', 2/))
2 Σ, 84
# t max (axx, k1)
# max maxx; 12)
yi 81/2
8; exp(3rr; 11)
8; exp (3x;; 11)
Return g
8 t exp(n' 14)
6 + exp(y' 11)
+ 52 + Es 8:
2 - * 2' + Ei
yi + s1/2
Return 5, (H, 8, z) Scaling (on prior values),
Return T. (11, 82',2)
1 = 7 8
Scaling (on prior values), 62'/2
Scaling (at tail),
of
ÿ 0/2
Table 4: We provide diagrams and code for various forms of SoftMax. , is the inverse tempuraturo paramcter, and is set to 0.5
Fusing SoftMax with a contraction limits the output size. The join tail allows it to be fused with a contraction. This limits the size of the output, yielding a streamable algorithm. To streamline the derivation, we do not explicitly draw the updated maintained variables as in Figure We can then apply rearrangements to recover a streamable form of the composed function.
Can replace with streamed contraction.
2 201
Transpose
moves back
30

---

Published in Transactions on Machine Learning Research (03/2025)
B
z/8/2
We can then replace SoftMax with unscaled SoftMax, which is done in FlashAAtention- This lets us move the shared "/=" factor to the end of the expression, producing the numerically stable form we use. We forego drawing the auxiliary variables lines to streamline the derivation but add them later.
8
/z
0
X
0
1
This analysis derives the streamability of SoftMax and its fusion with contraction as a standard procedure using diagrams. Diagrams, by showing the structure of constituent algorithms, act algebraic tools for deriving fusion.
B Performance Analysis
The performance analysis assisted by an Excel spreadsheet we developed, available al githabb.com//i mii- sardini-lab/ Nookki, In the future, additional tools will be provided that repository.
B.0.1 Lower Bound Compute Time
The lower bound on compute time is given by the clock cycles thrcad for tensor cores krC times the number of Uhreads In divided by the clock frequency fx. This must be larger than the transfer time, which is equal to the bytes per iteration H times the per SM bandwidth B B/NSM- We assume no caching and a sufficiently large F. These asscssments are given in Table 4 for different algorithms in the linked document. We artificially set " Num. of Th" 999 for the FlashAttention algorithms simulate not being memory bound by using caches.
31

---

Published in Transactions on Machine Learning Research (03/2025)
Tk 2 Tu KTC 9a/5k H B/NsM
97 2 JK H B/JK NsM)
3 295 (for H100 SXM5)
B.1 Arithmetic Intensity of Matrix Multiplication
With FLOPs 2abe, we are required to perform 0.5 + (2b)-1 transfers per FLOP, H/K. These computations and transfers occur across the GPU, mcaning we use GPU-wide bandwidth. Each compute requires transfer from GMEM to the L2 cache (at. 3352 GB/s) and transfer from the L2 cache to SMEM bandwidth of 1.7e12 and 6.0e12 FP16 values per second, respectively. As the L2 cache uses intermediate caching strategy, its effective memory is given by the total memory of child processors where data is ultimately stored (see Appendix A.2.1. Therefore, we use M, = NsM MsM-
(at 12 TB/s) (Shah et al. 2024 This translates
To avoid a bandwidth bottleneck, we require that M- + < We assemble these values and subsequent calculations in the attached document's "MatMul Bandwidth" sheet. The absolute mindmun values indicate the M and b values that, if not exceeded, will ensure the algorithm is bottlenecked, indepen- dent of other factors. The Min b values indicate the size of the inner-dimension must exceed, given the GPU memory, to not be bandwidth bottlenecked. Note that these expressions give mininums, larger and imperfect caching will slow down the algorithm.
Bandwidth Limits
FLOPs
Kt
9.894E 14
ops/s
Quantization (bytes/value)
q
B/value
Level
This corresponds to the caching level
Bandwidth (bytes)
3.35E+12
B/s
Bandwidth (values)
Ht L2
1.68E+12
values/s
Computes/Transfer
Kt/Ht
591
Absolute Min L2 Memory
q(Kt/Ht) 22
681
KB
Absolute Min b
(Kt/Ht)/2
295
values
N_ SM per L2
N. SM
66
Memory/Processor
M SM N SM
16896
KB
Tile Size
(0.5
2941
values
Min b
(0
370
values
SMEM Level
This corresponds to the lower level
Bandwidth (bytes)
1.20E 13
B/s
Bandwidth (values)
Ht SM
6.00E+12 values/s
Computes/Transfer
Kt/Ht
165
Absolute Min L2 Memory
q(Kt/Ht) *2
53 KB
Absolute Min b
(Kt/Ht)/2
82
values
Memory/Processor
M_SM
256 KB
Tile Size
(M/
362
values
Min b
(M/)
151
values
32

---

Published in Transactions on Machine Learning Research (03/2025)
B.2
Configuration Tables
B.2.1
Derived Attention
Table 1: Contiguration Table for Hopper Attention
Partitions
Per Warpgoup
(per q)
Extra Loaded
Replacement
Quantization [B]
Per Thread
Block
Variables
Level
d
d'
Multiple
Storage
Divisors
128
128
128
32
8
128
Axis Sizes
128
128
128
64
64
128
32
Q
Queries
Tensor Core
16384 SMEM
Keys
Values
Tensor Core
SMEM
S
Attention Scores
Tensor Core
16384
Registers
P
Weighted Scores
SMEM
16384
SMEM
P
Weighted Scores (Registers)
Registers
16384
Registers
A
Weighted Scores (Tensor Core)
Tensor Core
8192
SMEM
Auxiliary
Registers
3
768
Registers
O'
Output (egisters))
Registers
32768
Registers
D
Delta Output (Tensor Obre)
Tensor Core
1
8192
Registers
40
Delta Output (SMEM)
SMEM
8192
SMEM
40'
Delta Output (Registers)
Registers
2048
Registers
Total (Bytes)
49152
49152
SMEM
0
76544
Registers
Number of Warpgroups Configuration
N
384
Table 2 Maximum um Analysis
Table 4 Throuphput Analysis
Excess Calculations
H100 SXMS
H100 PCle Unit
Maximum Per TB
Per WG.
N max
Per TB
Per WG.
Per Th.
Clk/Th
SMEM
227
3.7
35
11.67
H/TB
24576
Registers
256
1.83-009
1,62E-09
Sources
Under CUDA Enabled center Products, Hopper CC9.0
Sandwitth
8
3.35E+12
2.04E-12
B/s
Technical Soncificattonn tions Compute Capability
Num. of SMs
N SM
152
114
Table 3 Compute Analysis
Bandwitth SM
B/N SM
254E+11
1.79E-10 B/s
Q-K
SoltMax
Load Time
H
SM/B
9.68E
1.37E-06
Type
FP16
Core
Function
g.
Ops/Th
2d
x/u 2
x/ux
Comp Time
TKK qg/1
1.25E-06
1.42E-06
From table
10384
65
16384
256
32768
63
256
Time
1 K. T.H
1,20E-08
1.42E-06
Ops/Ckk
8192
OpsЛloop
K=g.g g
120E+07
1.25E-07
4.0625
Throughput
SMNT
1.01E+15 FLOP/s
.39. (Oos/clk) IFL K
Table
CUDA Programming Guide, Arthhmelic Instructions
33

---

Published in Transactions on Machine Learning Research (03/2025)
B.2.2 FlashAttention-3 FP8 Inter- warpgrup
Table 1: Configuration Table for Hopper Attention
Partitions
Per Warpgroup oup
Extra Loaded
Replacement
Quantization
(per g g)
Per Thread
Block
Variables
Level
g
Muttiple
(B)
Storage
Divisors
128
32
128
Axis Sizes
128
128
Queries
Tensor Core
16384
Registers
K
Keys
Tensor Core
1
1
57344
SMEM
Values
Tensor Core
57344
SMEM
S
Scores Weights
Tensor Core
1
57344
Registers
Output
Tensor Core
16384
Registers Registers
a
Auxiliary
Registers
768
Total (Bytes)
114688
SMEM
90880
Registers
Excess Registers per Thread
314
Number of Warpgroups Confgurraton
g 256
Table 2: Maximum Analysis
Table 4: Throughput Analysis
SMEM
#DM/O!
H/TB
Registers
258
75
78.5
39.25
314
Cik/s
LK
1.89E+09
1,62E-00
Sources
Unneer CUDD& Enahled batacenter Products Hopper has 0C9.0
Bandwidth
8
3.35E+12
2.04E-12
B/s
Technical Specifications per Compute Capability
Num. of SMs
N.SM
132
114
Compute Analysis
1.79E-10
Q-K
SoftMax
P-V
Totals
Load Time
H HN. SM/B
2.29E-06
3.21E-06
Core Type
Tensor
Special
Tensor
Tensor Core
Special
Min. Th
T_H.K/K
295
371
Core
Function
Core
Function
Num. of Th
999
Ops/Th
2ds
2ds
Comp Time
T.K 8g g/9/1K
7.64E-06
8.63E-06 5
Tmmax(T
8.63E-06
Ops/Ckk
8192
16
8192
Ops/toop
"(K/g
1.19E+08
1.15E-08 FLOP
Oks/Th
14
14
14.00
Throughput
SM
196E+15
1.51E-15 FLOP/s
Sources
Table p39. (Ops/c /clk) (FLOPsMIN KM
Sources
Table, p39
CUDA Programming Guide, Artthmetic Instructions
B.2.3 FlashAttention-3 FP8 Intra-Warpgrout
Table Configuration Table for Hopper Attention
Partitions
Extra Loaded
Replacement
Quantization
Per Warpgroup p
(per q)
Per Thread
Block
Variables
Level
g
Multiple
(B)
Storage
Divisors
128
32
128
Sizes
128
128
Queries
Tensor Core
1
16384
Registers
K
Keys
Tensor Core
1
1
49152
SMEM
V
Values
Tensor Core
49152
SMEM
S
Scores Weights
Tensor Core
1
1
2
98304
Registers
Output
Tensor Core
1
16384
Registers
Auxiliary
Registers
1
2
768
Registers
Total (Bytes)
98304
SMEM
131 131840 840
Reglsters
Ercess Registers per Thread
34

---

Published in Transactions on Machine Learning Research (03/2025)
Number Warpgroups Configuration
N-
128
Table 2: Maximum Analysis
Table 4: Throughput Analysis
Excess Calculations
H100 SXMS H100 PCl Unit
Maximum Per TB.
Per WG.
N. max
Per TB.
PG
Per Tn.
Cik/Th
K
24
SMEM
227
96
#DW/O!
151
131.00
H/TB
H
40152
8
Registers
256
128.7 75
127.25
127.25
1018
Cik/s
LK
1.83E+09
1,625-08
Products. Hopper
2.04E-12
Technical Specifications Compute Capability
Num. of SMs
N.SM
114
Table 3: Compute Analysis
Bandwidth SM
B/N SM
2.54E+10
1.79E+10 B/s
Q-K
SoftMax
P-V
Totals
Load Time
H.H HN. SM/B
1.94E-06
2.75E-00 5
Core Type
Tensor
Special
Tensor
Tensor
Special
Min. Th
8-0 K.HKKKK
148
185
8.9
Ops/Th
2ds
x
Comp Time
Tk q/1K
1.31E-05
1.48E-05
From mbll
49152
192
49152
96304
192
Time
T-maxT K, K.T.K...... H)
1.31E-05
1.48E-05
Ops/Ckk
4096
16
4096
4036
16
Ops/toop
K=g.9 "KK/g
9.82E+07
9.82E-07 FLOP
Throughput
Table
(FLOPs)IN K)
Table,
OUDA Programming Guide, Arthmetic Instructions
B.2.4
FlashAttention-3 FP16 Intra- Parra-rarpgroup
Table 1: Configuration Table for Hopper Attention
Partitions
Extra Loaded
Replacement
Quantization
Per Warpgrous oup
(per g. g)
Per Thread
Block
Variables
Level
g
Multiple
(B)
Storage
Divisors
128
32
128
Axis Sizes
128
128
128
Queries
Tensor Core
1
32768
Registers
K
Keys
Tensor Core
1
1
1
2
65536
SMEM
Values
Tensor Core
1
65536
S
Scores Welghts
Tensor Core
1
1
1
2
65536
Reglsters
Output
Tensor Core
1
32768
Registers
Auxiliary
Registers
1
1536
Registers
Total (Bytes)
131072
SMEM
132608
Registers
Excess Registers per Thread
-12
Number Warggroups Configuration
N
g_ 128
Table 2: Maximum Analysis
Table 4: Throughhut Analysis
Excess Calculations
H100 SXMS H100 PCle Uhit
Maximum Per T8.
Per WG.
N max
Per T8.
WG
Per Th.
Cik/Th
16
SMEM
227
#DI//!
99.00 00
H/TB
65536
Registers
126.5
126.50
1012
Cik/s
1.83E-09
1.62E-00
Sources
Under QUDA Enabled Daracente enter Products. Hopper has CC9.0
Bandwidth
8
3.35E+12
2.04E+12
B/s
Technical Specilications per Compute Capability
Num. of SMs
N.SM
132
114
Table 3: Compute Analysis
Bandwidth SM
B/N
2.54+100
1.79E+10 B/s
SoftMax
Core
Tensor
Special
Tensor
Tensor Core
Special Function
Th
Core
Function
Core
Num. of
999
Ops/Th
2ds.x
s.x
2ds.x
Comp Time
8.73E-06
9.87E-06 5
From table
32768
32768
65536
128
Time
T= K,
8.73E-06
9.87E-06
K=g.9 ((K/g.))
6.55E+07
Oks/Th
16
8.00
Throughput
NM
9.89E+14
7.58E+14 FLOP/s
Sources
Table n30 (Ops//cl (clk) = (FLOPSMIN SMH K)
Sources
Table. 000
CUDA Programming Guide, Arithmelic Instructions
35


Reviews:
[{"review_id": "LdnoHD0UbV", "reviewer": "Reviewer_UB1n", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "### Strengths\n- The diagrammatic approach tries to simplify complex deep learning algorithm optimizations, making them more accessible and systematic.\n- The paper combines theoretical insights with practical hardware considerations, offering a holistic view of algorithm optimization for GPUs.\n\n### Weaknesses\n- The framework lacks experimental results to validate its theoretical performance improvements. At the very least, empirical validation under minimal settings should be provided for methods like FlashAttention-3 to substantiate the claims.\n- While diagrams are central to the contribution, some are dense and may overwhelm readers unfamiliar with the methodology. More intuitive explanations or detailed walkthroughs are needed.\n- The lack of guidelines for creating diagrams may lead to inconsistent representations across researchers, limiting the framework's utility.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "This paper presents a diagrammatic framework for optimizing deep learning algorithms, focusing on IO-awareness to address memory bandwidth and transfer bottlenecks in GPU-based computations. It proposes a performance model based on diagrams that depict data types, functions, and GPU hierarchies, allowing for systematic derivation of optimized algorithms. It also introduces techniques such as group partitioning and stream partitioning for resource-efficient computations. Additionally, the authors delve into hardware-specific optimizations, such as coalesced memory access and tensor-core operations, presenting pseudocode tailored for Ampere and Hopper architectures.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "- Please include experimental benchmarks comparing the proposed scheme with SOTA methods like FlashAttention-3. At a minimum, provide empirical validation under basic settings to demonstrate performance gains. Please report metrics such as execution time, and memory usage.\n- The process of understanding the calculations demonstrating that the proposed framework can fit up to 13 warps per thread block, as shown in Table 3, is overly complex and scattered across the appendix. It would be more effective to present the key computational steps clearly on the main pages for easier comprehension.\n- Please ensure consistent notation such as $$(\\mathbb{R}^{a \\times b \\times c} \\times \\mathbb{R}^{d \\times e \\times c}) \\rightarrow (\\mathbb{R}^{a \\times b \\times c} \\times \\mathbb{R}^{d \\times e}) $$ in Section 2.1.", "broader_impact_concerns": "N/A", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "PBICbf8xjJ", "reviewer": "Reviewer_JqHN", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "This is a highly unusual technical paper with loads of details, and I can only imagine the immense amount of work put into such paper. The idea of using diagrams to visualize and analyze algorithms is intriguing and could potentially lead to better understanding of problems and further enhancing the solutions.\n\nWeaknesses (Not all of the following are necessarily weaknesses, but rather things that are not entirely clear to me and areas of confusion that could be addressed or clarified):\n\nFirst, I find myself struggling with the paper's claim that it can be \"used to quickly derive methods like FlashAttention.\" To my best understsanding of the paper, the approach seems primarily post-hoc - going from known algorithm -> diagram representation -> analysis, rather than going from problem -> diagram analysis -> discover new algorithm. While the analysis is thorough, I'm not convinced that someone unfamiliar with FlashAttention could discover similar optimizations using just this diagrammatic method.\n\nThis raises a more fundamental question: what is the value proposition if we already understand algorithms like FlashAttention? The method appears capable of suggesting theoretical optimizations - for instance, fitting 13 warps per SM instead of FlashAttention's 4-8. It seems that it may be capable of suggesting theoretical optimizations to further optimize approaches that are already known. But without empirical validation, it's hard to judge whether these optimizations would work in practice. I think the paper lacks empirical rigor despite being highly technical and information dense.\n\nThe learning curve also concerns me. The framework requires understanding both the diagrammatic notation and hardware concepts. I wonder if the complexity is justified by the analytical insights gained. While I can see its utility as an analytical tool for IO-aware optimization, I'm less convinced about its broader adoption given the investment required to use it effectively.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "The paper provides a diagrammatic framework, based on Neural Circuit Diagrams for representing and analyzing algorithms focusing on hardware constraints and optimization. Through this method it proposes a theoretical optimization to FlashAttention that suggests potential for improved hardware utilization, though it remains untested.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "1. I would like to see more clarity about what this framework can and cannot do. If it's primarily an analytical tool rather than a method for discovering new algorithms, the paper should be more precise in its claims.\n\n2. The paper would benefit from comparing it with existing performance modeling tools (if any exist). Why should researchers invest time in learning such method? Some concrete comparisons would help make the case.\n\n3. The paper is quite dense and challenging to follow. More intuitive examples and clearer explanations of hardware concepts would help readers without extensive hardware backgrounds.\n\n4. Most importantly, I think some empirical validation would greatly strengthen the paper. Is it feasible to implement and test some of their proposed optimizations? Concrete benchmarks would make the theoretical contributions more convincing. This paper is already difficult to follow with a steep learning curve, but demonstrating its practicality and effectiveness could convince engineers and researchers to adopt this method.", "broader_impact_concerns": "No concerns.", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "qC6KGzvxXI", "reviewer": "Reviewer_rQeH", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "**strengths**\n- the paper is very well written, the approach is introduced step by step, each section building on the definitions and insights of the previous one, towards progressively more complex algorithms\n- the topic of diagrammatic representation of DL algorithms can be of interest to at least part of the audience of this journal, and prompt further work in this area\n- to the best of my knowledge, the proposed approach is novel and unique\n- the authors show that thanks to their framework they are able to extract useful insights from existing algorithms and craft hardware-specific alternatives with the potential to outperform existing strategies\n\n**weaknesses**\n- the main drawback is the lack of on-hardware verification of the theoretical estimates for memory consumption and transfer costs derived using the proposed framework, as well as a lack of demonstration of the efficacy of the newly derived attention algorithms\n\n**other comments**\n- in Figure 12: what is the factor $N_g$ appearing in the cumulative per-group transfer cost $H_{l1,g}$ formula?\n- Figure 15 is not referenced in the text\n- in Section 3.1 and Figure 18, could you spell out the derivation of the memory consumption inequalities for matmul ($\\sqrt M \\ge g_c$) and why \"the transfer cost of (...) matrix multiplication is cubic for $n \\ge \\sqrt M/2$\". Also, why is it assumed that a = b = c, isn't this a loss of generality on the shape of matrices?\n- Section 4.3 deals with values expressed with fewer bits, and comes to the natural conclusion that quantization is beneficial for throughput (everything else being equal). However, quantization algorithms typically rely on storing additional quantization parameters (such as scaling factor) that are used to quantize / dequantize the matrix. This represent an overhead in terms of memory consumption, transfer costs, as well as operations. For example, for a matrix of size N x M and group size GS, group-quantization stores (N / GS) x M scales, twice as many if asymmetric quantization is enabled. It would be interesting expanding this section to consider the whole impact of some actual quantization algorithms, like the aforementioned per-group. As presented, the finding of this section appears to be trivial", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "This paper introduces a diagrammatic representation of Deep Learning (DL) algorithms. It defines rules and building blocks of this approach, and connects them to implement operations and algorithms of various complexity, offering estimates of memory usage (across memory levels) and transfer costs. Such algorithms are then expanded into pseudo-code with explicit representation of loops and variables, allowing for a direct comparison with existing algorithms like FlashAttention. Using the proposed approach, the authors derive two attention algorithms, for Ampere and Hopper architectures, which may theoretically outperform conventional attention algorithms.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "I have a positive view of this paper in its present form but I would appreciate if the authors could address my comments in the other section.", "broader_impact_concerns": "No concern on my end.", "additional_comments": "", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Meta-review: Overall this is an interesting and comprehensive work introducing a novel framework for analyzing algorithms running on a GPU and their performance characteristics. All reviewers appreciated this novelty and the technical contribution of this paper, as well as the use of FlashAttention to illustrate the framework's potential utility. While it would have been nice to provide empirical validation to really demonstrate the utility of the framework, as it is, the contribution of the diagrammatic analysis framework already satisfies TMLR's criteria for acceptance.

Please revise the paper accordingly to update the references (e.g. (?) on Page 18, including publication years, and venues instead of arXiv for published works).
Decision: Accept (poster)